diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/state_fusion.py b/src/gt4py/next/program_processors/runners/dace/transformations/state_fusion.py index 0e1f24feb9..af460c194b 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/state_fusion.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/state_fusion.py @@ -163,6 +163,16 @@ def _check_for_wcr_conflicts( if len(common_write_data) == 0: return False + # Transients are assumed to be written only once (ADR-18), in a single + # state. If a transient is written to in both states — e.g. partial + # writes on disjoint subsets, as produced for `concat_where` — merging + # the states would leave one of the writer AccessNodes unconnected to + # any reader (the consumer rewiring in `apply` only handles + # second-state source nodes), leading to dangling dead writes or + # uninitialized reads. Reject the fusion in this case. + if any(sdfg.arrays[data].transient for data in common_write_data): + return True + for state, ac_nodes in [ (first_state, first_state_writes), (second_state, second_state_writes), diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_state_fusion.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_state_fusion.py index ad6f6841f5..ca3d1635d7 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_state_fusion.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_state_fusion.py @@ -172,6 +172,68 @@ def _make_global_both_state_write() -> tuple[dace.SDFG, dace.SDFGState, dace.SDF return sdfg, state1, state2 +def _make_transient_both_states_write() -> tuple[dace.SDFG, dace.SDFGState, dace.SDFGState]: + """In both states the same transient is written and it is read in the second one. + + The two states write disjoint subsets of the transient `t` and the second + state also fully reads it. This is, in essence, the pattern that is emitted + by the stree lowering of `concat_where`. + + Note that the fusion is not allowed, because after it the write to `t` in + the first state would no longer be connected to the read in the second + state (ADR-018 requires that a transient is written within a single state). + """ + sdfg = dace.SDFG(util.unique_name("transient_write_in_both_states")) + state1 = sdfg.add_state(is_start_block=True) + state2 = sdfg.add_state_after(state1) + + for name in ["a", "b", "o"]: + sdfg.add_array( + name, + shape=(20,), + dtype=dace.float64, + transient=False, + ) + sdfg.add_array( + "t", + shape=(20,), + dtype=dace.float64, + transient=True, + ) + + state1.add_mapped_tasklet( + "comp1", + map_ranges={"__i": "0:10"}, + inputs={"__in1": dace.Memlet("a[__i]")}, + code="__out = __in1 + 10.", + outputs={"__out": dace.Memlet("t[__i]")}, + external_edges=True, + ) + + t2 = state2.add_access("t") + state2.add_mapped_tasklet( + "comp2_1", + map_ranges={"__i": "10:20"}, + inputs={"__in1": dace.Memlet("b[__i]")}, + code="__out = __in1 + 12.", + outputs={"__out": dace.Memlet("t[__i]")}, + external_edges=True, + output_nodes={t2}, + ) + state2.add_mapped_tasklet( + "comp2_2", + map_ranges={"__i": "0:20"}, + inputs={"__in1": dace.Memlet("t[__i]")}, + code="__out = __in1 * 2.", + outputs={"__out": dace.Memlet("o[__i]")}, + external_edges=True, + input_nodes={t2}, + ) + sdfg.validate() + + return sdfg, state1, state2 + + def _make_empty_state( first_state_empty: bool, ) -> tuple[dace.SDFG, dace.SDFGState, dace.SDFGState]: @@ -589,6 +651,22 @@ def test_global_in_both_states_write(): ) +def test_transient_in_both_states_write(): + sdfg, state1, state2 = _make_transient_both_states_write() + + # The transient is written in both states and read in the second one. Fusing + # the states would disconnect the write to `t` in the first state from the + # read in the second state, so the transformation must not apply. + nb_applied = sdfg.apply_transformations_repeated(gtx_transformations.GT4PyStateFusion) + sdfg.validate() + + assert nb_applied == 0 + assert sdfg.start_block is state1 + assert sdfg.number_of_nodes() == 2 + assert util.count_nodes(state1, dace_nodes.AccessNode) == 2 + assert util.count_nodes(state2, dace_nodes.AccessNode) == 3 + + def test_empty_first_state(): sdfg, state1, state2 = _make_empty_state(True) assert state1.number_of_nodes() == 0