From e17f0f423d3663fd6b974dc4754af0a233639d44 Mon Sep 17 00:00:00 2001 From: Lukas Truemper Date: Sun, 26 Nov 2023 10:33:07 +0100 Subject: [PATCH] Copy2Map: formatted --- dace/transformation/dataflow/copy_to_map.py | 10 ++++++++-- tests/transformations/copy_to_map_test.py | 6 +++++- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/dace/transformation/dataflow/copy_to_map.py b/dace/transformation/dataflow/copy_to_map.py index 195add15f8..24532780b7 100644 --- a/dace/transformation/dataflow/copy_to_map.py +++ b/dace/transformation/dataflow/copy_to_map.py @@ -108,11 +108,17 @@ def apply(self, state: SDFGState, sdfg: SDFG): # Linearize and delinearize to get index expression for other side if copy_a: - a_index = [symbolic.pystr_to_symbolic(f'__i{i} + {edge.data.src_subset.ranges[i][0]}') for i in range(len(copy_shape))] + a_index = [ + symbolic.pystr_to_symbolic(f'__i{i} + {edge.data.src_subset.ranges[i][0]}') + for i in range(len(copy_shape)) + ] b_index = self.delinearize_linearize(bdesc, copy_shape, edge.data.get_dst_subset(edge, state)) else: a_index = self.delinearize_linearize(adesc, copy_shape, edge.data.get_src_subset(edge, state)) - b_index = [symbolic.pystr_to_symbolic(f'__i{i} + {edge.data.dst_subset.ranges[i][0]}') for i in range(len(copy_shape))] + b_index = [ + symbolic.pystr_to_symbolic(f'__i{i} + {edge.data.dst_subset.ranges[i][0]}') + for i in range(len(copy_shape)) + ] a_subset = subsets.Range([(ind, ind, 1) for ind in a_index]) b_subset = subsets.Range([(ind, ind, 1) for ind in b_index]) diff --git a/tests/transformations/copy_to_map_test.py b/tests/transformations/copy_to_map_test.py index b340bec496..6d6004b8fd 100644 --- a/tests/transformations/copy_to_map_test.py +++ b/tests/transformations/copy_to_map_test.py @@ -7,6 +7,7 @@ def _copy_to_map(storage: dace.StorageType): + @dace def somecopy(a, b): b[:] = a @@ -24,6 +25,7 @@ def somecopy(a, b): def _flatten_to_map(storage: dace.StorageType): + @dace def somecopy(a, b): b[:] = a.flatten() @@ -103,6 +105,7 @@ def test_preprocess(): def test_squeezed_subset_src(): + @dace.program def squeezed_subset_src(a: dace.float32[10], b: dace.float32[10, 10]): for i in range(10): @@ -126,10 +129,11 @@ def squeezed_subset_src(a: dace.float32[10], b: dace.float32[10, 10]): def test_squeezed_subset_src_swapped(): + @dace.program def squeezed_subset_src_swapped(a: dace.float32[10], b: dace.float32[10, 10]): for i in range(10): - b[i+1:10, i] = a[0:10-(i+1)] + b[i + 1:10, i] = a[0:10 - (i + 1)] sdfg = squeezed_subset_src_swapped.to_sdfg()