Fix: DDE removing read from access_set in read/write nodes

[romanc]

Description

For this bug to show, we need two separate states with a transient produced in one and subsequently read and written (but not read again). It is important that StateFusion isn't able to merge these two state. I've put a dummy if/else in the middle. Before DDE this might look like

where tmp_computed is transient and tmp is a given variable. DDE will now go and see that tmp_computed can be removed as an output of the read_write tasklet. The currently faulty update of access_set will remove tmp_computed from the list of reads in block state. This will then propagate (badly) up to the start state where tmp_computed is marked as never read again, removing the whole tasklet, leaving the block state to read an uninitialized tmp_computed (if we were to codegen).

Repro

import dace
import os

# Create an empty SDFG
sdfg = dace.SDFG(os.path.basename(__file__).removesuffix(".py").replace("-", "_"))

sdfg.add_scalar("tmp", dace.float32)
sdfg.add_scalar("tmp_computed", dace.float32, transient=True)

start_state = sdfg.add_state("start", is_start_block=True)

read_tmp = start_state.add_read("tmp")
write_computed = start_state.add_write("tmp_computed")

# upstream tasklet that writes a transient (to be read in a separate state)
write = start_state.add_tasklet("write", {"IN_tmp"}, {"OUT_computed"}, "OUT_computed = IN_tmp * 2 + 1")
start_state.add_memlet_path(read_tmp, write, dst_conn="IN_tmp", memlet=dace.Memlet(data="tmp"))
start_state.add_memlet_path(write, write_computed, src_conn="OUT_computed", memlet=dace.Memlet(data="tmp_computed"))

# Add a condition to avoid fusing next_state and start_state separate
guard = sdfg.add_state_after(start_state, "guard_state")
true_state = sdfg.add_state("true_state")
false_state = sdfg.add_state("false_state")
fs_read = false_state.add_read("tmp")
fs_write = false_state.add_write("tmp")
fs_tasklet = false_state.add_tasklet("abs", {"IN_tmp"}, {"OUT_tmp"}, "OUT_tmp = -IN_tmp")
false_state.add_memlet_path(fs_read, fs_tasklet, dst_conn="IN_tmp", memlet=dace.Memlet("tmp"))
false_state.add_memlet_path(fs_tasklet, fs_write, src_conn="OUT_tmp", memlet=dace.Memlet("tmp"))
merge = sdfg.add_state("merge_state")

sdfg.add_edge(guard, true_state, dace.InterstateEdge("tmp >= 0"))
sdfg.add_edge(guard, false_state, dace.InterstateEdge("tmp < 0"))
sdfg.add_edge(true_state, merge, dace.InterstateEdge())
sdfg.add_edge(false_state, merge, dace.InterstateEdge())

next_state = sdfg.add_state_after(merge)

write_computed = next_state.add_write("tmp_computed")
read_computed = next_state.add_read("tmp_computed")
write_tmp = next_state.add_write("tmp")

# downstream tasklet that reads _and_ writes a transient
consume = next_state.add_tasklet("read_write", {"IN_computed"}, {"OUT_tmp", "OUT_computed"}, "OUT_computed = 2 * IN_computed\nOUT_tmp = OUT_computed + IN_computed")
next_state.add_memlet_path(read_computed, consume, dst_conn="IN_computed", memlet=dace.Memlet(data="tmp_computed"))
next_state.add_memlet_path(consume, write_tmp, src_conn="OUT_tmp", memlet=dace.Memlet(data="tmp"))
next_state.add_memlet_path(consume, write_computed, src_conn="OUT_computed", memlet=dace.Memlet(data="tmp_computed"))

sdfg.validate()

sdfg.simplify(verbose=True, validate=True)

assert len(list(filter(lambda node: isinstance(node, dace.sdfg.nodes.Tasklet), sdfg.start_block.nodes()))) == 1, "write tasklet in start_block is gone"

Finishing this PR

I'll need help to evaluate whether or not the proposed solution is a good one. Questions that I have:

• Is this the right place to fix it?
• Should we manually change access_sets or would it be simpler/more reliable to redo the analysis step?
• The repro case translates to a unit test. Is it a good one or should I e.g. call DDE directly and write assertions for the output of the pass?

[tbennun]

Good find! Looks like a bug in DDE indeed.

I need to review these remaining access nodes a bit more, but so far looks like a good solution.

If you can create a pipeline that calls DeadDataflowElimination that would be better. You can check out the rest of the DDE tests in the relevant file (tests/passes/dead_dataflow_elimination_test.py) to see how the rest of the tests invoke DDE.

[romanc]

Just running DDE instead of the full pipeline was a good idea. It allowed me to simplify the test case and narrowly test for the expected cleanup.

simplified test case

[romanc]

PS: the same issue also seem present in main. Branch https://github.com/romanc/dace/tree/romanc/dde-fix-mainline ports this fix to mainline. Let's have all the discussions here first and then once we all agree, the mainline version is just gonna be a port.