From 171ddcae8becfd81cc5cace9846bdd40e0be1b6a Mon Sep 17 00:00:00 2001 From: Tal Ben-Nun Date: Tue, 5 Sep 2023 15:20:37 -0700 Subject: [PATCH 1/3] Enable more arguments for `with dace.tasklet` --- dace/frontend/python/astutils.py | 14 ++++++++++++++ dace/frontend/python/interface.py | 3 ++- dace/frontend/python/newast.py | 17 +++++++++++++++++ dace/frontend/python/preprocessing.py | 2 +- 4 files changed, 34 insertions(+), 2 deletions(-) diff --git a/dace/frontend/python/astutils.py b/dace/frontend/python/astutils.py index 4a0ec88531..faf214fdeb 100644 --- a/dace/frontend/python/astutils.py +++ b/dace/frontend/python/astutils.py @@ -705,3 +705,17 @@ def escape_string(value: Union[bytes, str]): return value.encode("unicode_escape").decode("utf-8") # Python 2.x return value.encode('string_escape') + + +def parse_function_arguments(node: ast.Call, argnames: List[str]) -> Dict[str, ast.AST]: + """ + Parses function arguments (both positional and keyword) from a Call node, + based on the function's argument names. If an argument was not given, it will + not be in the result. + """ + result = {} + for arg, aname in zip(node.args, argnames): + result[aname] = arg + for kw in node.keywords: + result[kw.arg] = kw.value + return result diff --git a/dace/frontend/python/interface.py b/dace/frontend/python/interface.py index ea1970dafd..69e650beaa 100644 --- a/dace/frontend/python/interface.py +++ b/dace/frontend/python/interface.py @@ -293,10 +293,11 @@ class tasklet(metaclass=TaskletMetaclass): The DaCe framework cannot analyze these tasklets for optimization. """ - def __init__(self, language: Union[str, dtypes.Language] = dtypes.Language.Python): + def __init__(self, language: Union[str, dtypes.Language] = dtypes.Language.Python, side_effects: bool = False): if isinstance(language, str): language = dtypes.Language[language] self.language = language + self.side_effects = side_effects def __enter__(self): if self.language != dtypes.Language.Python: diff --git a/dace/frontend/python/newast.py b/dace/frontend/python/newast.py index c9d92b7860..b5d27e14f4 100644 --- a/dace/frontend/python/newast.py +++ b/dace/frontend/python/newast.py @@ -2510,6 +2510,7 @@ def _parse_tasklet(self, state: SDFGState, node: TaskletType, name=None): # Looking for the first argument in a tasklet annotation: @dace.tasklet(STRING HERE) langInf = None + side_effects = None if isinstance(node, ast.FunctionDef) and \ hasattr(node, 'decorator_list') and \ isinstance(node.decorator_list, list) and \ @@ -2522,6 +2523,19 @@ def _parse_tasklet(self, state: SDFGState, node: TaskletType, name=None): langArg = node.decorator_list[0].args[0].value langInf = dtypes.Language[langArg] + # Extract arguments from with statement + if isinstance(node, ast.With): + expr = node.items[0].context_expr + if isinstance(expr, ast.Call): + args = astutils.parse_function_arguments(expr, ['language', 'side_effects']) + langArg = args.get('language', None) + side_effects = args.get('side_effects', None) + langInf = astutils.evalnode(langArg, {**self.globals, **self.defined}) + if isinstance(langInf, str): + langInf = dtypes.Language[langInf] + + side_effects = astutils.evalnode(side_effects, {**self.globals, **self.defined}) + ttrans = TaskletTransformer(self, self.defined, self.sdfg, @@ -2536,6 +2550,9 @@ def _parse_tasklet(self, state: SDFGState, node: TaskletType, name=None): symbols=self.symbols) node, inputs, outputs, self.accesses = ttrans.parse_tasklet(node, name) + if side_effects is not None: + node.side_effects = side_effects + # Convert memlets to their actual data nodes for i in inputs.values(): if not isinstance(i, tuple) and i.data in self.scope_vars.keys(): diff --git a/dace/frontend/python/preprocessing.py b/dace/frontend/python/preprocessing.py index 10a1ab120e..239875118f 100644 --- a/dace/frontend/python/preprocessing.py +++ b/dace/frontend/python/preprocessing.py @@ -1268,7 +1268,7 @@ def _convert_to_ast(contents: Any): node) else: # Augment closure with new value - newnode = self.resolver.global_value_to_node(e, node, f'inlined_{id(contents)}', True, keep_object=True) + newnode = self.resolver.global_value_to_node(contents, node, f'inlined_{id(contents)}', True, keep_object=True) return newnode return _convert_to_ast(contents) From 2adc2e565fa8ea7012ecb0b87e9ab9dc42d1d12c Mon Sep 17 00:00:00 2001 From: Tal Ben-Nun Date: Tue, 5 Sep 2023 15:21:03 -0700 Subject: [PATCH 2/3] More informative message when using explicit tasklets with wrong dimensionality --- dace/frontend/python/memlet_parser.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/dace/frontend/python/memlet_parser.py b/dace/frontend/python/memlet_parser.py index 6ef627a430..7cc218c4fb 100644 --- a/dace/frontend/python/memlet_parser.py +++ b/dace/frontend/python/memlet_parser.py @@ -285,7 +285,11 @@ def ParseMemlet(visitor, if len(node.value.args) >= 2: write_conflict_resolution = node.value.args[1] - subset, new_axes, arrdims = parse_memlet_subset(array, node, das, parsed_slice) + try: + subset, new_axes, arrdims = parse_memlet_subset(array, node, das, parsed_slice) + except IndexError: + raise DaceSyntaxError(visitor, node, 'Failed to parse memlet expression due to dimensionality. ' + f'Array dimensions: {array.shape}, expression in code: {astutils.unparse(node)}') # If undefined, default number of accesses is the slice size if num_accesses is None: From 4b98f0cdfa631538113acbd9271d6b32c15c30ab Mon Sep 17 00:00:00 2001 From: Tal Ben-Nun Date: Tue, 5 Sep 2023 21:46:52 -0700 Subject: [PATCH 3/3] Fix test --- tests/numpy/advanced_indexing_test.py | 477 +++++++++++++------------- 1 file changed, 246 insertions(+), 231 deletions(-) diff --git a/tests/numpy/advanced_indexing_test.py b/tests/numpy/advanced_indexing_test.py index 48853cdf26..d2c348ce95 100644 --- a/tests/numpy/advanced_indexing_test.py +++ b/tests/numpy/advanced_indexing_test.py @@ -1,231 +1,246 @@ -# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -""" -Tests for numpy advanced indexing syntax. See also: -https://numpy.org/devdocs/reference/arrays.indexing.html -""" -import dace -import numpy as np -import pytest - -N = dace.symbol('N') -M = dace.symbol('M') - - -def test_flat(): - @dace.program - def indexing_test(A: dace.float64[20, 30]): - return A.flat - - A = np.random.rand(20, 30) - res = indexing_test(A) - assert np.allclose(A.flat, res) - - -def test_flat_noncontiguous(): - with dace.config.set_temporary('compiler', 'allow_view_arguments', value=True): - - @dace.program - def indexing_test(A): - return A.flat - - A = np.random.rand(20, 30).transpose() - res = indexing_test(A) - assert np.allclose(A.flat, res) - - -def test_ellipsis(): - @dace.program - def indexing_test(A: dace.float64[5, 5, 5, 5, 5]): - return A[1:5, ..., 0] - - A = np.random.rand(5, 5, 5, 5, 5) - res = indexing_test(A) - assert np.allclose(A[1:5, ..., 0], res) - - -def test_aug_implicit(): - @dace.program - def indexing_test(A: dace.float64[5, 5, 5, 5, 5]): - A[:, 1:5][:, 0:2] += 5 - - A = np.random.rand(5, 5, 5, 5, 5) - regression = np.copy(A) - regression[:, 1:5][:, 0:2] += 5 - indexing_test(A) - assert np.allclose(A, regression) - - -def test_ellipsis_aug(): - @dace.program - def indexing_test(A: dace.float64[5, 5, 5, 5, 5]): - A[1:5, ..., 0] += 5 - - A = np.random.rand(5, 5, 5, 5, 5) - regression = np.copy(A) - regression[1:5, ..., 0] += 5 - indexing_test(A) - assert np.allclose(A, regression) - - -def test_newaxis(): - @dace.program - def indexing_test(A: dace.float64[20, 30]): - return A[:, np.newaxis, None, :] - - A = np.random.rand(20, 30) - res = indexing_test(A) - assert res.shape == (20, 1, 1, 30) - assert np.allclose(A[:, np.newaxis, None, :], res) - - -def test_multiple_newaxis(): - @dace.program - def indexing_test(A: dace.float64[10, 20, 30]): - return A[np.newaxis, :, np.newaxis, np.newaxis, :, np.newaxis, :, np.newaxis] - - A = np.random.rand(10, 20, 30) - res = indexing_test(A) - assert res.shape == (1, 10, 1, 1, 20, 1, 30, 1) - assert np.allclose(A[np.newaxis, :, np.newaxis, np.newaxis, :, np.newaxis, :, np.newaxis], res) - - -def test_index_intarr_1d(): - @dace.program - def indexing_test(A: dace.float64[N], indices: dace.int32[M]): - return A[indices] - - A = np.random.rand(20) - indices = [1, 10, 15] - res = indexing_test(A, indices, M=3) - assert np.allclose(A[indices], res) - - -def test_index_intarr_1d_literal(): - @dace.program - def indexing_test(A: dace.float64[20]): - return A[[1, 10, 15]] - - A = np.random.rand(20) - indices = [1, 10, 15] - res = indexing_test(A) - assert np.allclose(A[indices], res) - - -def test_index_intarr_1d_constant(): - indices = [1, 10, 15] - - @dace.program - def indexing_test(A: dace.float64[20]): - return A[indices] - - A = np.random.rand(20) - res = indexing_test(A) - assert np.allclose(A[indices], res) - - -def test_index_intarr_1d_multi(): - @dace.program - def indexing_test(A: dace.float64[20, 10, 30], indices: dace.int32[3]): - return A[indices, 2:7:2, [15, 10, 1]] - - A = np.random.rand(20, 10, 30) - indices = [1, 10, 15] - res = indexing_test(A, indices) - # FIXME: NumPy behavior is unclear in this case - assert np.allclose(np.diag(A[indices, 2:7:2, [15, 10, 1]]), res) - - -def test_index_intarr_nd(): - @dace.program - def indexing_test(A: dace.float64[4, 3], rows: dace.int64[2, 2], columns: dace.int64[2, 2]): - return A[rows, columns] - - A = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]], dtype=np.float64) - rows = np.array([[0, 0], [3, 3]], dtype=np.intp) - columns = np.array([[0, 2], [0, 2]], dtype=np.intp) - expected = A[rows, columns] - res = indexing_test(A, rows, columns) - assert np.allclose(expected, res) - - -def test_index_boolarr_rhs(): - @dace.program - def indexing_test(A: dace.float64[20, 30]): - return A[A > 15] - - A = np.ndarray((20, 30), dtype=np.float64) - for i in range(20): - A[i, :] = np.arange(0, 30) - regression = A[A > 15] - - # Right-hand side boolean array indexing is unsupported - with pytest.raises(IndexError): - res = indexing_test(A) - assert np.allclose(regression, res) - - -def test_index_multiboolarr(): - @dace.program - def indexing_test(A: dace.float64[20, 20], B: dace.bool[20]): - A[B, B] = 2 - - A = np.ndarray((20, 20), dtype=np.float64) - for i in range(20): - A[i, :] = np.arange(0, 20) - B = A[:, 1] > 0 - - # Advanced indexing with multiple boolean arrays should be disallowed - with pytest.raises(IndexError): - indexing_test(A, B) - - -def test_index_boolarr_fixed(): - @dace.program - def indexing_test(A: dace.float64[20, 30], barr: dace.bool[20, 30]): - A[barr] += 5 - - A = np.ndarray((20, 30), dtype=np.float64) - for i in range(20): - A[i, :] = np.arange(0, 30) - barr = A > 15 - regression = np.copy(A) - regression[barr] += 5 - - indexing_test(A, barr) - - assert np.allclose(regression, A) - - -def test_index_boolarr_inline(): - @dace.program - def indexing_test(A: dace.float64[20, 30]): - A[A > 15] = 2 - - A = np.ndarray((20, 30), dtype=np.float64) - for i in range(20): - A[i, :] = np.arange(0, 30) - regression = np.copy(A) - regression[A > 15] = 2 - - indexing_test(A) - - assert np.allclose(regression, A) - - -if __name__ == '__main__': - test_flat() - test_flat_noncontiguous() - test_ellipsis() - test_aug_implicit() - test_ellipsis_aug() - test_newaxis() - test_multiple_newaxis() - test_index_intarr_1d() - test_index_intarr_1d_literal() - test_index_intarr_1d_constant() - test_index_intarr_1d_multi() - test_index_intarr_nd() - test_index_boolarr_rhs() - test_index_multiboolarr() - test_index_boolarr_fixed() - test_index_boolarr_inline() +# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. +""" +Tests for numpy advanced indexing syntax. See also: +https://numpy.org/devdocs/reference/arrays.indexing.html +""" +import dace +from dace.frontend.python.common import DaceSyntaxError +import numpy as np +import pytest + +N = dace.symbol('N') +M = dace.symbol('M') + + +def test_flat(): + + @dace.program + def indexing_test(A: dace.float64[20, 30]): + return A.flat + + A = np.random.rand(20, 30) + res = indexing_test(A) + assert np.allclose(A.flat, res) + + +def test_flat_noncontiguous(): + with dace.config.set_temporary('compiler', 'allow_view_arguments', value=True): + + @dace.program + def indexing_test(A): + return A.flat + + A = np.random.rand(20, 30).transpose() + res = indexing_test(A) + assert np.allclose(A.flat, res) + + +def test_ellipsis(): + + @dace.program + def indexing_test(A: dace.float64[5, 5, 5, 5, 5]): + return A[1:5, ..., 0] + + A = np.random.rand(5, 5, 5, 5, 5) + res = indexing_test(A) + assert np.allclose(A[1:5, ..., 0], res) + + +def test_aug_implicit(): + + @dace.program + def indexing_test(A: dace.float64[5, 5, 5, 5, 5]): + A[:, 1:5][:, 0:2] += 5 + + A = np.random.rand(5, 5, 5, 5, 5) + regression = np.copy(A) + regression[:, 1:5][:, 0:2] += 5 + indexing_test(A) + assert np.allclose(A, regression) + + +def test_ellipsis_aug(): + + @dace.program + def indexing_test(A: dace.float64[5, 5, 5, 5, 5]): + A[1:5, ..., 0] += 5 + + A = np.random.rand(5, 5, 5, 5, 5) + regression = np.copy(A) + regression[1:5, ..., 0] += 5 + indexing_test(A) + assert np.allclose(A, regression) + + +def test_newaxis(): + + @dace.program + def indexing_test(A: dace.float64[20, 30]): + return A[:, np.newaxis, None, :] + + A = np.random.rand(20, 30) + res = indexing_test(A) + assert res.shape == (20, 1, 1, 30) + assert np.allclose(A[:, np.newaxis, None, :], res) + + +def test_multiple_newaxis(): + + @dace.program + def indexing_test(A: dace.float64[10, 20, 30]): + return A[np.newaxis, :, np.newaxis, np.newaxis, :, np.newaxis, :, np.newaxis] + + A = np.random.rand(10, 20, 30) + res = indexing_test(A) + assert res.shape == (1, 10, 1, 1, 20, 1, 30, 1) + assert np.allclose(A[np.newaxis, :, np.newaxis, np.newaxis, :, np.newaxis, :, np.newaxis], res) + + +def test_index_intarr_1d(): + + @dace.program + def indexing_test(A: dace.float64[N], indices: dace.int32[M]): + return A[indices] + + A = np.random.rand(20) + indices = [1, 10, 15] + res = indexing_test(A, indices, M=3) + assert np.allclose(A[indices], res) + + +def test_index_intarr_1d_literal(): + + @dace.program + def indexing_test(A: dace.float64[20]): + return A[[1, 10, 15]] + + A = np.random.rand(20) + indices = [1, 10, 15] + res = indexing_test(A) + assert np.allclose(A[indices], res) + + +def test_index_intarr_1d_constant(): + indices = [1, 10, 15] + + @dace.program + def indexing_test(A: dace.float64[20]): + return A[indices] + + A = np.random.rand(20) + res = indexing_test(A) + assert np.allclose(A[indices], res) + + +def test_index_intarr_1d_multi(): + + @dace.program + def indexing_test(A: dace.float64[20, 10, 30], indices: dace.int32[3]): + return A[indices, 2:7:2, [15, 10, 1]] + + A = np.random.rand(20, 10, 30) + indices = [1, 10, 15] + res = indexing_test(A, indices) + # FIXME: NumPy behavior is unclear in this case + assert np.allclose(np.diag(A[indices, 2:7:2, [15, 10, 1]]), res) + + +def test_index_intarr_nd(): + + @dace.program + def indexing_test(A: dace.float64[4, 3], rows: dace.int64[2, 2], columns: dace.int64[2, 2]): + return A[rows, columns] + + A = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]], dtype=np.float64) + rows = np.array([[0, 0], [3, 3]], dtype=np.intp) + columns = np.array([[0, 2], [0, 2]], dtype=np.intp) + expected = A[rows, columns] + res = indexing_test(A, rows, columns) + assert np.allclose(expected, res) + + +def test_index_boolarr_rhs(): + + @dace.program + def indexing_test(A: dace.float64[20, 30]): + return A[A > 15] + + A = np.ndarray((20, 30), dtype=np.float64) + for i in range(20): + A[i, :] = np.arange(0, 30) + regression = A[A > 15] + + # Right-hand side boolean array indexing is unsupported + with pytest.raises(IndexError): + res = indexing_test(A) + assert np.allclose(regression, res) + + +def test_index_multiboolarr(): + + @dace.program + def indexing_test(A: dace.float64[20, 20], B: dace.bool[20]): + A[B, B] = 2 + + A = np.ndarray((20, 20), dtype=np.float64) + for i in range(20): + A[i, :] = np.arange(0, 20) + B = A[:, 1] > 0 + + # Advanced indexing with multiple boolean arrays should be disallowed + with pytest.raises(DaceSyntaxError): + indexing_test(A, B) + + +def test_index_boolarr_fixed(): + + @dace.program + def indexing_test(A: dace.float64[20, 30], barr: dace.bool[20, 30]): + A[barr] += 5 + + A = np.ndarray((20, 30), dtype=np.float64) + for i in range(20): + A[i, :] = np.arange(0, 30) + barr = A > 15 + regression = np.copy(A) + regression[barr] += 5 + + indexing_test(A, barr) + + assert np.allclose(regression, A) + + +def test_index_boolarr_inline(): + + @dace.program + def indexing_test(A: dace.float64[20, 30]): + A[A > 15] = 2 + + A = np.ndarray((20, 30), dtype=np.float64) + for i in range(20): + A[i, :] = np.arange(0, 30) + regression = np.copy(A) + regression[A > 15] = 2 + + indexing_test(A) + + assert np.allclose(regression, A) + + +if __name__ == '__main__': + test_flat() + test_flat_noncontiguous() + test_ellipsis() + test_aug_implicit() + test_ellipsis_aug() + test_newaxis() + test_multiple_newaxis() + test_index_intarr_1d() + test_index_intarr_1d_literal() + test_index_intarr_1d_constant() + test_index_intarr_1d_multi() + test_index_intarr_nd() + test_index_boolarr_rhs() + test_index_multiboolarr() + test_index_boolarr_fixed() + test_index_boolarr_inline()