fix[next-dace]: block state fusion creating a WAR hazard on global data

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

@@ -14,6 +14,10 @@
 from dace.sdfg import nodes as dace_nodes, utils as dace_sdutils
 from ordered_set import OrderedSet
 
+from gt4py.next.program_processors.runners.dace.transformations import (
+    utils as gtx_transformation_utils,
+)
+
 
 # Conditional import because `gt4py.cartesian` uses an older DaCe version without
 #  `explicit_cf_compatible`.
@@ -307,6 +311,16 @@ def _check_for_read_write_dependencies(
                     #  dataflow is executed. We thus have to ensure that the new
                     #  component will not affect any component that is not involved
                     #  inside the merge.
+                    # Moreover, the consumer itself must not create a WAR hazard:
+                    #  the messenger only places the producer's writes before the
+                    #  consumer's accesses, it should not place the producer's *reads*
+                    #  of the same data before the consumer's write. Thus, we refuse
+                    #  such cases.
+                    if any(
+                        self._has_war_hazard(first_state, second_state, data, consumed_messenger)
+                        for data in producer_dependency.intersection(consumer_influece)
+                    ):
+                        return True
                     unaffected_producers: list[int] = [
                         producer_id
                         for producer_id, data_producer in enumerate(data_producers)
@@ -337,6 +351,90 @@ def _check_for_read_write_dependencies(
 
         return False
 
+    def _has_war_hazard(
+        self,
+        first_state: dace.SDFGState,
+        second_state: dace.SDFGState,
+        data: str,
+        messengers: set[str],
+    ) -> bool:
+        """Checks if the fusion would create a write-after-read hazard on `data`.
+
+        The data `data` is read in the first state and written to in the second
+        state, by a consumer component that consumes the `messengers` from the
+        first state. In the merged state the consumer's write is only ordered after
+        the data the write depends on, i.e. the messengers. Thus the write is
+        unordered with respect to a read in the first state - a WAR hazard - unless
+        that read is dataflow-upstream of the exchanged messages, in which case the
+        dataflow through the messenger imposes the order.
+
+        The check is performed at node level and considers the fusion safe for a
+        write of `data` in the second state if there is an exchanged messenger such
+        that:
+        - the messenger is dataflow-upstream of the write in the second state, and
+        - every read of `data` in the first state is dataflow-upstream of the
+          messenger's producing AccessNode in the first state.
+
+        Args:
+            first_state: The first state of the fusion.
+            second_state: The second state of the fusion.
+            data: The global data that is read in the first state and written to
+                in the second state.
+            messengers: The data that is produced by a component of the first state
+                and consumed by the consumer component of the second state.
+
+        Returns:
+            `True` if the fusion would create a WAR hazard on `data`.
+        """
+
+        # The sites where `data` is read in the first state are identified by the
+        #  destination nodes of the read edges, i.e. the MapEntries and Tasklets
+        #  that perform the read.
+        read_sites = {
+            edge.dst
+            for read_node in first_state.data_nodes()
+            if read_node.data == data and first_state.out_degree(read_node) != 0
+            for edge in first_state.out_edges(read_node)
+        }
+        write_nodes = [
+            node
+            for node in second_state.data_nodes()
+            if node.data == data and second_state.in_degree(node) != 0
+        ]
+        for write_node in write_nodes:
+            upstream_of_write = gtx_transformation_utils.find_upstream_nodes(
+                write_node, second_state
+            )
+            for messenger in messengers:
+                # The AccessNodes that produce the messenger in the first state.
+                messenger_producers = [
+                    node
+                    for node in first_state.data_nodes()
+                    if node.data == messenger and first_state.in_degree(node) != 0
+                ]
+                # The AccessNodes through which the messenger enters the consumer
+                #  component of the second state.
+                has_war_path = any(
+                    node.data == messenger and second_state.out_degree(node) != 0
+                    for node in upstream_of_write
+                    if isinstance(node, dace_nodes.AccessNode)
+                )
+
+                if has_war_path and all(
+                    any(
+                        gtx_transformation_utils.is_reachable(read_site, producer, first_state)
+                        for producer in messenger_producers
+                    )
+                    for read_site in read_sites
+                ):
+                    # The write is ordered after all reads through `messenger`.
+                    break
+
+            else:  # belongs to `for messenger in messengers`
+                return True
+
+        return False
+
     def _move_nodes(self, sdfg: dace.SDFG) -> None:
         """Moves the nodes from the second state to the first state.
 

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

@@ -780,3 +780,106 @@ def test_hidden_double_producer():
     sdfg.validate()
 
     assert nb_applied == 0
+
+
+def _make_global_read_up_and_downstream_of_messenger() -> tuple[
+    dace.SDFG, dace.SDFGState, dace.SDFGState
+]:
+    """The second state writes a global that the first state also reads downstream of the messenger.
+
+    The first state reads the global `g` twice; once to produce the transient
+    messenger `t` (`t = g + 1`) and once together with the messenger
+    (`u = t / g`). Note that both computations share the same AccessNodes and
+    therefore form a single concurrent subgraph. The second state then performs
+    the write back `g = t`.
+
+    This models in-place update field operators of the form
+
+    ```python
+    z_g = g
+    g_new = g + incr
+    o = g_new / z_g
+    ```
+
+    where the second state performs the write back of `g_new` into `g`.
+    """
+    sdfg = dace.SDFG(util.unique_name("global_read_up_and_downstream_of_messenger"))
+    state1 = sdfg.add_state(is_start_block=True)
+    state2 = sdfg.add_state_after(state1)
+
+    for name in ["g", "t", "u"]:
+        sdfg.add_array(
+            name,
+            shape=(10,),
+            dtype=dace.float64,
+            transient=name in {"t", "u"},
+        )
+
+    # The AccessNodes are shared between the two computations, such that the
+    #  dataflow of the first state forms a single concurrent subgraph.
+    g = state1.add_access("g")
+    t = state1.add_access("t")
+
+    # Upstream computation: `t = g + 1`, `t` is the messenger.
+    state1.add_mapped_tasklet(
+        "comp1",
+        map_ranges={"__i": "0:10"},
+        inputs={"__in": dace.Memlet("g[__i]")},
+        code="__out = __in + 1.",
+        outputs={"__out": dace.Memlet("t[__i]")},
+        input_nodes={g},
+        output_nodes={t},
+        external_edges=True,
+    )
+    # Downstream computation: `u = t / g`, i.e. `g` is read a second time,
+    #  downstream of the messenger `t`.
+    state1.add_mapped_tasklet(
+        "comp2",
+        map_ranges={"__i": "0:10"},
+        inputs={
+            "__in1": dace.Memlet("t[__i]"),
+            "__in2": dace.Memlet("g[__i]"),
+        },
+        code="__out = __in1 / __in2",
+        outputs={"__out": dace.Memlet("u[__i]")},
+        input_nodes={t, g},
+        external_edges=True,
+    )
+
+    # Second state: write back `g = t`.
+    state2.add_nedge(
+        state2.add_access("t"),
+        state2.add_access("g"),
+        dace.Memlet("t[0:10] -> [0:10]"),
+    )
+
+    return sdfg, state1, state2
+
+
+def test_global_read_up_and_downstream_of_messenger():
+    """Fusing the states would create a WAR hazard on `g`, so it must not happen.
+
+    In the merged state the write `g = t` is dataflow-ordered after the
+    computation of the messenger `t`, and thus after the *first* read of `g`.
+    However, it is not ordered with respect to the *second* read of `g`: the
+    computation `u = t / g` depends on the messenger as well and is therefore
+    concurrent with the write back.
+    """
+    sdfg, state1, state2 = _make_global_read_up_and_downstream_of_messenger()
+    assert sdfg.number_of_nodes() == 2
+    # Both computations of the first state form a single concurrent subgraph.
+    assert len(dace.sdfg.utils.concurrent_subgraphs(state1)) == 1
+    assert util.count_nodes(state1, dace_nodes.AccessNode) == 3
+
+    nb_applied = sdfg.apply_transformations_repeated(gtx_transformations.GT4PyStateFusion)
+    sdfg.validate()
+
+    assert nb_applied == 0
+    assert sdfg.number_of_nodes() == 2
+
+    # Ensure that the unfused SDFG computes the intended result, i.e. the write
+    #  back `g = t` happened only after all reads of `g` in the first state.
+    g_orig = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0])
+    g = g_orig.copy()
+    sdfg(g=g)
+    assert np.allclose(g, g_orig + 1.0)