From 1a2085326a0410b65e721bc2f1d4790514b1350e Mon Sep 17 00:00:00 2001 From: Tal Ben-Nun Date: Wed, 22 Mar 2023 18:40:51 -0700 Subject: [PATCH 1/2] Fix internal subscript access if already existed --- dace/frontend/python/newast.py | 6 ++-- tests/numpy/subarray_in_nested_call_test.py | 31 ++++++++++++++++++++- 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/dace/frontend/python/newast.py b/dace/frontend/python/newast.py index 9789a433e2..7e99b1b19a 100644 --- a/dace/frontend/python/newast.py +++ b/dace/frontend/python/newast.py @@ -3021,9 +3021,9 @@ def _add_read_access(self, if name in self.sdfg.arrays: return (name, None) elif (name, rng, 'w') in self.accesses: - return self.accesses[(name, rng, 'w')] + return self.accesses[(name, rng, 'w')][0], rng elif (name, rng, 'r') in self.accesses: - return self.accesses[(name, rng, 'r')] + return self.accesses[(name, rng, 'r')][0], rng elif name in self.variables: return (self.variables[name], None) elif name in self.scope_vars: @@ -3047,7 +3047,7 @@ def _add_write_access(self, if name in self.sdfg.arrays: return (name, None) if (name, rng, 'w') in self.accesses: - return self.accesses[(name, rng, 'w')] + return self.accesses[(name, rng, 'w')][0], rng elif name in self.variables: return (self.variables[name], None) elif (name, rng, 'r') in self.accesses or name in self.scope_vars: diff --git a/tests/numpy/subarray_in_nested_call_test.py b/tests/numpy/subarray_in_nested_call_test.py index 09227b3f83..322c660da0 100644 --- a/tests/numpy/subarray_in_nested_call_test.py +++ b/tests/numpy/subarray_in_nested_call_test.py @@ -1,4 +1,5 @@ -# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +import math import numpy as np import dace @@ -45,6 +46,34 @@ def test_inout_connector(): assert np.allclose(a, ref) +def test_indirect_symbolic_access(): + + @dace.program + def tester(a: dace.float64[20], b: dace.float64[3], c: dace.float64[19]): + for i in dace.map[0:10]: + xtmp: dace.float64 = 0 + + local_offset_i = (i + 1) % 2 + + if local_offset_i < 3 % 2: + for kx in dace.unroll(range(math.ceil(3 / 2))): + ind_i = (i + 1 - kx * 2) // 2 + kind_i = local_offset_i + kx * 2 + if ind_i >= 0 and ind_i < 20: + xtmp += a[ind_i] * b[kind_i] + + c[i] = xtmp + + a = np.random.rand(20) + b = np.random.rand(15) + c = np.random.rand(10) + refc = np.copy(c) + tester.f(a, b, refc) + tester(a, b, c) + assert np.allclose(c, refc) + + if __name__ == '__main__': test() test_inout_connector() + test_indirect_symbolic_access() From 3ad07c05cbc5bc7871efb2e394417bbc366e93cf Mon Sep 17 00:00:00 2001 From: Tal Ben-Nun Date: Thu, 23 Mar 2023 10:57:38 -0700 Subject: [PATCH 2/2] Better fix and another test --- dace/frontend/python/newast.py | 24 ++++++++++++---------- tests/python_frontend/indirections_test.py | 19 +++++++++++++++++ 2 files changed, 32 insertions(+), 11 deletions(-) diff --git a/dace/frontend/python/newast.py b/dace/frontend/python/newast.py index 7e99b1b19a..d9a6458bf8 100644 --- a/dace/frontend/python/newast.py +++ b/dace/frontend/python/newast.py @@ -3020,23 +3020,25 @@ def _add_read_access(self, arr_type: data.Data = None): if name in self.sdfg.arrays: return (name, None) - elif (name, rng, 'w') in self.accesses: - return self.accesses[(name, rng, 'w')][0], rng - elif (name, rng, 'r') in self.accesses: - return self.accesses[(name, rng, 'r')][0], rng elif name in self.variables: return (self.variables[name], None) + + if (name, rng, 'w') in self.accesses: + new_name, new_rng = self.accesses[(name, rng, 'w')] + elif (name, rng, 'r') in self.accesses: + new_name, new_rng = self.accesses[(name, rng, 'r')] elif name in self.scope_vars: new_name, new_rng = self._add_access(name, rng, 'r', target, new_name, arr_type) - full_rng = subsets.Range.from_array(self.sdfg.arrays[new_name]) - if (_subset_has_indirection(rng, self) or _subset_is_local_symbol_dependent(rng, self)): - new_name, new_rng = self.make_slice(new_name, rng) - elif full_rng != new_rng: - new_name, new_rng = self.make_slice(new_name, new_rng) - return (new_name, new_rng) else: raise NotImplementedError + full_rng = subsets.Range.from_array(self.sdfg.arrays[new_name]) + if (_subset_has_indirection(rng, self) or _subset_is_local_symbol_dependent(rng, self)): + new_name, new_rng = self.make_slice(new_name, rng) + elif full_rng != new_rng: + new_name, new_rng = self.make_slice(new_name, new_rng) + return (new_name, new_rng) + def _add_write_access(self, name: str, rng: subsets.Range, @@ -3047,7 +3049,7 @@ def _add_write_access(self, if name in self.sdfg.arrays: return (name, None) if (name, rng, 'w') in self.accesses: - return self.accesses[(name, rng, 'w')][0], rng + return self.accesses[(name, rng, 'w')] elif name in self.variables: return (self.variables[name], None) elif (name, rng, 'r') in self.accesses or name in self.scope_vars: diff --git a/tests/python_frontend/indirections_test.py b/tests/python_frontend/indirections_test.py index 0ea6cf5cf0..c59dffb922 100644 --- a/tests/python_frontend/indirections_test.py +++ b/tests/python_frontend/indirections_test.py @@ -59,6 +59,23 @@ def test_indirection_scalar_nsdfg(): assert (np.allclose(res, A[x])) +@dc.program +def indirection_scalar2_nsdfg(A: dc.float64[10], x: dc.int32[10]): + B = np.empty_like(A) + for i in dc.map[0:A.shape[0]]: + a = x[i] + B[i] = A[a] + B[i] = A[a] + return B + + +def test_indirection_scalar2_nsdfg(): + A = np.random.randn(10).astype(np.float64) + x = np.random.randint(0, 10, size=(10, ), dtype=np.int32) + res = indirection_scalar2_nsdfg(A, x) + assert (np.allclose(res, A[x])) + + @dc.program def indirection_scalar_assign_nsdfg(A: dc.float64[10], x: dc.int32[10]): B = np.empty_like(A) @@ -169,6 +186,7 @@ def test_indirection_scalar_range(): def test_indirection_scalar_range_nsdfg(): + @dc.program def indirection_scalar_range_nsdfg(A: dc.float64[10], x: dc.int32[11]): B = np.empty_like(A) @@ -374,6 +392,7 @@ def test_spmv(): test_indirection_scalar_assign() test_indirection_scalar_augassign() test_indirection_scalar_nsdfg() + test_indirection_scalar2_nsdfg() test_indirection_scalar_assign_nsdfg() test_indirection_scalar_augassign_nsdfg() test_indirection_scalar_multi()