Fix codegen with data access on inter-state edge (#1434)

After uplift to dace v0.15, one SDFG which was working before started to
show compilation errors. The latest DaCe is moving a data access to an
inter-state edge. For the data-access, the symbols that define array
strides are needed for code generation. The SDFG was validated, before
and after the simplify pass, but it did not compile for CPU. When
skipping the simplify pass, the compilation did work. The problem has
been narrowed down to the scalar-to-symbol promotion, which is moving a
data access to an inter-state edge. Then, the method
`_used_symbols_internal` needs to be update to account for data
containers, including symbolic shape and strides.

This issue was reported in #1433. This PR contains a unit test to
reproduce the issue and verify the proposed fix.

dace/sdfg/state.py

@@ -2493,6 +2493,10 @@ def _used_symbols_internal(self,
                 # subracting the (true) free symbols from the edge's assignment keys. This way we can correctly
                 # compute the symbols that are used before being assigned.
                 efsyms = e.data.used_symbols(all_symbols)
+                # collect symbols representing data containers
+                dsyms = {sym for sym in efsyms if sym in self.arrays}
+                for d in dsyms:
+                    efsyms |= {str(sym) for sym in self.arrays[d].used_symbols(all_symbols)}
                 defined_syms |= set(e.data.assignments.keys()) - (efsyms | state_symbols)
                 used_before_assignment.update(efsyms - defined_syms)
                 free_syms |= efsyms

tests/codegen/codegen_used_symbols_test.py

@@ -88,8 +88,55 @@ def test_codegen_used_symbols_gpu():
         pass
 
 
+def test_codegen_edge_assignment_with_indirection():
+    rng = numpy.random.default_rng(42)
+    (M, N, K) = (dace.symbol(x, dace.int32) for x in ['M', 'N', 'K'])
+
+    sdfg = dace.SDFG('edge_assignment_with_indirection')
+    [sdfg.add_symbol(x, dace.int32) for x in {'__indirect_idx', '__neighbor_idx'}]
+    sdfg.add_array('_field', (M,), dace.float64)
+    sdfg.add_array('_table', (N,K), dace.int32)
+    sdfg.add_array('_out', (N,), dace.float64)
+
+    state0 = sdfg.add_state(is_start_block=True)
+    state1 = sdfg.add_state()
+    sdfg.add_edge(state0, state1, dace.InterstateEdge(
+        assignments={'_field_idx': '_table[__indirect_idx, __neighbor_idx]'}
+    ))
+    state1.add_memlet_path(
+        state1.add_access('_field'),
+        state1.add_access('_out'),
+        memlet=dace.Memlet(data='_out', subset='__indirect_idx', other_subset='_field_idx', wcr='lambda x, y: x + y')
+    )
+
+    M, N, K = (5, 4, 2)
+    field = rng.random((M,))
+    out = rng.random((N,))
+    table = numpy.random.randint(0, M, (N, K), numpy.int32)
+
+    TEST_INDIRECT_IDX = numpy.random.randint(0, N)
+    TEST_NEIGHBOR_IDX = numpy.random.randint(0, K)
+
+    reference = numpy.asarray(
+        [
+            out[i] + field[table[i, TEST_NEIGHBOR_IDX]] if i == TEST_INDIRECT_IDX else out[i]
+            for i in range(N)
+        ]
+    )
+
+    sdfg(
+        _field=field, _table=table, _out=out, M=M, N=N, K=K,
+        __indirect_idx=TEST_INDIRECT_IDX,
+        __neighbor_idx=TEST_NEIGHBOR_IDX
+    )
+
+    assert numpy.allclose(out, reference)
+
+
 if __name__ == "__main__":
 
     test_codegen_used_symbols_cpu()
     test_codegen_used_symbols_cpu_2()
     test_codegen_used_symbols_gpu()
+    test_codegen_edge_assignment_with_indirection()
+