diff --git a/dace/sdfg/state.py b/dace/sdfg/state.py index 461b18b1a9..becebd1c28 100644 --- a/dace/sdfg/state.py +++ b/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 diff --git a/tests/codegen/codegen_used_symbols_test.py b/tests/codegen/codegen_used_symbols_test.py index afa0ca0a05..1e216e9508 100644 --- a/tests/codegen/codegen_used_symbols_test.py +++ b/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() +