fix[next-dace]: do not relocate distributed buffer write back into loop bodies

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

@@ -531,6 +531,11 @@ class DistributedBufferRelocator(dace_transformation.Pass):
     - There is a `dest_storage` access node, that has an output degree larger
         than one.
 
+    Furthermore, the relocation will not happen if the state where `temp_storage`
+    is defined is nested inside a loop region. There the write back would be
+    executed in every iteration, potentially on data that the loop has only
+    written partially.
+
     Note:
         - Essentially this transformation removes the double buffering of
             `dest_storage`. Because we ensure that that `dest_storage` is non
@@ -681,6 +686,11 @@ def find_upstream_states(dst_state: dace.SDFGState) -> set[dace.SDFGState]:
             temp_storage_node, temp_storage_state = temp_storage
             def_locations: list[AccessLocation] = []
             for upstream_state in find_upstream_states(temp_storage_state):
+                if self._is_inside_loop(sdfg, upstream_state):
+                    # The definition of `temp_storage` is inside a loop. Moving the
+                    #  write back there would execute it in every iteration, thus
+                    #  on data that has potentially only been written partially.
+                    continue
                 if self._is_written_to_in_state(
                     data=temp_storage_node.data,
                     state=upstream_state,
@@ -764,6 +774,28 @@ def find_upstream_states(dst_state: dace.SDFGState) -> set[dace.SDFGState]:
 
         return result
 
+    def _is_inside_loop(
+        self,
+        sdfg: dace.SDFG,
+        state: dace.SDFGState,
+    ) -> bool:
+        """Checks if `state` is located inside a loop region.
+
+        Args:
+            sdfg: The SDFG on which we operate.
+            state: The state that should be examined.
+
+        Returns:
+            `True` if `state` is nested inside a `LoopRegion`, at any depth,
+            `False` otherwise.
+        """
+        scope = state.parent_graph
+        while scope is not None and scope is not sdfg:
+            if isinstance(scope, dace.sdfg.state.LoopRegion):
+                return True
+            scope = scope.parent_graph
+        return False
+
     def _is_written_to_in_state(
         self,
         data: str,

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

@@ -386,3 +386,62 @@ def test_distributed_buffer_conditional_block():
 
     res = gtx_transformations.gt_reduce_distributed_buffering(sdfg)
     assert res[sdfg]["DistributedBufferRelocator"][wb_state] == {"t"}
+
+
+def _make_distributed_buffer_definition_in_loop_sdfg() -> tuple[
+    dace.SDFG, dace.SDFGState, dace.SDFGState
+]:
+    """Creates an SDFG where the temporary is written inside a loop body.
+
+    The SDFG models what is generated for `scan` field operators: in iteration
+    `k` the loop writes element `k` of the temporary `t` and afterwards, in a
+    regular state, the fully written `t` is copied into the output `b`.
+    """
+    sdfg = dace.SDFG(util.unique_name("distributed_buffer_definition_in_loop_sdfg"))
+
+    for name in ["a", "b", "t"]:
+        sdfg.add_array(name=name, shape=(10,), dtype=dace.float64, transient=False)
+    sdfg.arrays["t"].transient = True
+    sdfg.add_symbol("k", dace.int32)
+
+    entry_state = sdfg.add_state(is_start_block=True)
+    loop_region = dace.sdfg.state.LoopRegion("loop", "k < 10", "k", "k = k + 1")
+    sdfg.add_node(loop_region)
+    sdfg.add_edge(entry_state, loop_region, dace.InterstateEdge(assignments={"k": 0}))
+
+    body_state = loop_region.add_state("loop_body", is_start_block=True)
+    scan_step = body_state.add_tasklet(
+        name="scan_step",
+        inputs={"__in"},
+        code="__out = __in + 1.0",
+        outputs={"__out"},
+    )
+    body_state.add_edge(body_state.add_access("a"), None, scan_step, "__in", dace.Memlet("a[k]"))
+    body_state.add_edge(scan_step, "__out", body_state.add_access("t"), None, dace.Memlet("t[k]"))
+
+    wb_state = sdfg.add_state_after(loop_region)
+    wb_state.add_nedge(wb_state.add_access("t"), wb_state.add_access("b"), dace.Memlet("t[0:10]"))
+
+    sdfg.validate()
+    return sdfg, body_state, wb_state
+
+
+def test_distributed_buffer_definition_in_loop():
+    """Tests that the write back is never moved into the body of a loop.
+
+    Relocating the write back into the loop would execute it in every iteration,
+    thus on data that the loop has only written partially, which is not only
+    useless, but also creates a data race.
+    """
+    sdfg, body_state, wb_state = _make_distributed_buffer_definition_in_loop_sdfg()
+    assert wb_state.number_of_nodes() == 2
+    assert not any(dnode.data == "b" for dnode in body_state.data_nodes())
+
+    res = gtx_transformations.gt_reduce_distributed_buffering(sdfg)
+
+    # The write back has to stay in the state after the loop; in particular no
+    #  write to `b` may be added inside the loop.
+    assert res is None or "DistributedBufferRelocator" not in res[sdfg]
+    assert wb_state.number_of_nodes() == 2
+    assert not any(dnode.data == "b" for dnode in body_state.data_nodes())
+    sdfg.validate()