From e83bec359c7655df6164c8e7afd709b18ecde6c6 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Mon, 24 Aug 2026 15:05:01 +0200 Subject: [PATCH] fix[next-dace]: Do not relocate distributed buffer write back into loop bodies --- .../runners/dace/transformations/simplify.py | 32 ++++++++++ .../test_distributed_buffer_relocator.py | 59 +++++++++++++++++++ 2 files changed, 91 insertions(+) diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/simplify.py b/src/gt4py/next/program_processors/runners/dace/transformations/simplify.py index 0c62972a6c..8418f9e111 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/simplify.py +++ b/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, diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_distributed_buffer_relocator.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_distributed_buffer_relocator.py index 6fcef3cf3e..b2f2e7f3d4 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_distributed_buffer_relocator.py +++ b/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()