From 1c0b0f6b810fa2792951d6a61bd3044915caee36 Mon Sep 17 00:00:00 2001 From: edopao Date: Tue, 5 Dec 2023 11:05:45 +0100 Subject: [PATCH] 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 | 4 ++ tests/codegen/codegen_used_symbols_test.py | 47 ++++++++++++++++++++++ 2 files changed, 51 insertions(+) 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() +