Copy2Map: Support for non-trivial subsets

dace/transformation/dataflow/copy_to_map.py

@@ -1,5 +1,6 @@
 # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved.
 
+import copy
 from dace import dtypes, symbolic, data, subsets, Memlet
 from dace.sdfg.scope import is_devicelevel_gpu
 from dace.transformation import transformation as xf
@@ -34,6 +35,11 @@ def can_be_applied(self, graph: SDFGState, expr_index: int, sdfg: SDFG, permissi
         if self.a.desc(sdfg).strides == self.b.desc(sdfg).strides:
             return False
 
+        # Happens-before edges
+        for edge in graph.edges_between(self.a, self.b):
+            if edge.data.data is None:
+                return False
+
         return True
 
     def delinearize_linearize(self, desc: data.Array, copy_shape: Tuple[symbolic.SymbolicType],
@@ -49,7 +55,32 @@ def delinearize_linearize(self, desc: data.Array, copy_shape: Tuple[symbolic.Sym
             return subsets.Range([(ind, ind, 1) for ind in indices])
 
         if rng is not None:  # Deal with offsets and strides in range
-            indices = rng.coord_at(indices)
+            # Artificial / unsqueezed dimensions
+            if len(rng) < len(copy_shape):
+                j = 0
+                unsqueeze_dims = []
+                for i in range(len(copy_shape)):
+                    if j >= len(desc.shape):
+                        unsqueeze_dims.append(i)
+                        continue
+
+                    copy_dim = copy_shape[i]
+                    if not ((isinstance(copy_dim, int) or copy_dim.is_Number) and copy_dim == 1):
+                        j += 1
+                        continue
+
+                    desc_dim = desc.shape[j]
+                    if ((isinstance(desc_dim, int) or desc_dim.is_Number) and desc_dim == 1):
+                        j += 1
+                        continue
+
+                    unsqueeze_dims.append(i)
+
+                unsqueezed_rng = copy.deepcopy(rng)
+                unsqueezed_rng.unsqueeze(unsqueeze_dims)
+                indices = unsqueezed_rng.coord_at(indices)
+            else:
+                indices = rng.coord_at(indices)
 
         linear_index = sum(indices[i] * data._prod(copy_shape[i + 1:]) for i in range(len(indices)))
 
@@ -82,11 +113,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}') 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}') 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])

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()
@@ -102,9 +104,84 @@ def test_preprocess():
     assert np.allclose(out, inp)
 
 
+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):
+            b[i, 0:i + 1] = a[0:i + 1]
+
+    sdfg = squeezed_subset_src.to_sdfg()
+
+    A = np.random.random(10).astype(np.float32)
+    B = np.random.random((10, 10)).astype(np.float32)
+    A_ = np.copy(A)
+    B_ = np.copy(B)
+
+    sdfg(A, B)
+
+    assert sdfg.apply_transformations(CopyToMap) == 1
+
+    sdfg(A_, B_)
+
+    assert np.allclose(A_, A)
+    assert np.allclose(B_, B)
+
+
+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)]
+
+    sdfg = squeezed_subset_src_swapped.to_sdfg()
+
+    A = np.random.random(10).astype(np.float32)
+    B = np.random.random((10, 10)).astype(np.float32)
+    A_ = np.copy(A)
+    B_ = np.copy(B)
+
+    sdfg(A, B)
+
+    assert sdfg.apply_transformations(CopyToMap) == 1
+
+    sdfg(A_, B_)
+
+    assert np.allclose(A_, A)
+    assert np.allclose(B_, B)
+
+
+def test_squeezed_subset_dst():
+
+    @dace.program
+    def squeezed_subset_dst(A: dace.float32[10], B: dace.float32[10, 10]):
+        for i in range(10):
+            A[0:10] = B[i, 0:10]
+
+    sdfg = squeezed_subset_dst.to_sdfg()
+
+    A = np.random.random(10).astype(np.float32)
+    B = np.random.random((10, 10)).astype(np.float32)
+    A_ = np.copy(A)
+    B_ = np.copy(B)
+
+    sdfg(A=A, B=B)
+
+    assert sdfg.apply_transformations(CopyToMap) == 1
+
+    sdfg(A_, B_)
+
+    assert np.allclose(A_, A)
+    assert np.allclose(B_, B)
+
+
 if __name__ == '__main__':
     test_copy_to_map()
     test_copy_to_map_gpu()
     test_flatten_to_map()
     test_flatten_to_map_gpu()
     test_preprocess()
+    test_squeezed_subset_src()
+    test_squeezed_subset_src_swapped()
+    test_squeezed_subset_dst()