Backport of #2165 (#2166)

Co-authored-by: Tal Ben-Nun <tbennun@users.noreply.github.com>

Author: phschaad

dace/frontend/python/newast.py

@@ -26,7 +26,7 @@
 from dace.frontend.python import nested_call, replacements, preprocessing
 from dace.frontend.python.memlet_parser import DaceSyntaxError, parse_memlet, ParseMemlet, inner_eval_ast, MemletExpr
 from dace.sdfg import nodes
-from dace.sdfg.propagation import propagate_memlet, propagate_subset, propagate_states
+from dace.sdfg.propagation import propagate_memlet, propagate_subset, propagate_states, align_memlet
 from dace.memlet import Memlet
 from dace.properties import LambdaProperty, CodeBlock
 from dace.sdfg import SDFG, SDFGState
@@ -2774,7 +2774,12 @@ def _add_assignment(self,
                     memlet.other_subset = op_subset
                     if op:
                         memlet.wcr = LambdaProperty.from_string('lambda x, y: x {} y'.format(op))
-                    state.add_nedge(op1, op2, memlet)
+                    if isinstance(self.sdfg.arrays[target_name], data.Reference):
+                        e = state.add_edge(op1, None, op2, 'set', memlet)
+                        # Align memlet to referenced array
+                        e.data = align_memlet(state, e, dst=False)
+                    else:
+                        state.add_nedge(op1, op2, memlet)
             else:
                 memlet = Memlet("{a}[{s}]".format(a=target_name,
                                                   s=','.join(['__i%d' % i for i in range(len(target_subset))])))
@@ -3272,9 +3277,10 @@ def visit_AnnAssign(self, node: ast.AnnAssign):
             storage = dtypes.StorageType.Default
             type_name = rname(node.annotation)
             warnings.warn('typeclass {} is not supported'.format(type_name))
-        if node.value is None and dtype is not None:  # Annotating type without assignment
+        if dtype is not None:
             self.annotated_types[rname(node.target)] = dtype
-            return
+            if node.value is None:  # Annotating type without assignment
+                return
         results = self._visit_assign(node, node.target, None, dtype=dtype)
         if storage != dtypes.StorageType.Default:
             self.sdfg.arrays[results[0][0]].storage = storage
@@ -3403,6 +3409,12 @@ def _visit_assign(self, node, node_target, op, dtype=None, is_return=False):
                         true_name, new_data = self.sdfg.add_temp_transient([1], result_data.dtype)
                         self.variables[name] = true_name
                         defined_vars[name] = true_name
+                    elif name in self.annotated_types and isinstance(self.annotated_types[name], data.Reference):
+                        desc = copy.deepcopy(self.annotated_types[name])
+                        desc.transient = True
+                        true_name = self.sdfg.add_datadesc(name, desc, find_new_name=True)
+                        self.variables[name] = true_name
+                        defined_vars[name] = true_name
                     elif (not name.startswith('__return')
                           and (isinstance(result_data, data.View) or
                                (not result_data.transient and isinstance(result_data, data.Array)))):

tests/sdfg/reference_test.py

@@ -10,6 +10,70 @@
 import networkx as nx
 
 
+def test_frontend_reference():
+    N = dace.symbol('N')
+    M = dace.symbol('M')
+    mystruct = dace.data.Structure(members={
+        "data": dace.data.Array(dace.float32, (N, M), strides=(1, N)),
+        "arrA": dace.data.ArrayReference(dace.float32, (N, )),
+        "arrB": dace.data.ArrayReference(dace.float32, (N, )),
+    },
+                                   name="MyStruct")
+
+    @dace.program
+    def init_prog(mydat: mystruct, fill_value: int) -> None:
+        mydat.arrA = mydat.data[:, 2]
+        mydat.arrB = mydat.data[:, 0]
+
+        # loop over all arrays and initialize them with `fill_value`
+        for index in range(M):
+            mydat.data[:, index] = fill_value
+
+        # Initialize the two named ones by name
+        mydat.arrA[:] = fill_value + 1
+        mydat.arrB[:] = fill_value + 2
+
+    dat = np.zeros((10, 5), dtype=np.float32)
+    inp_struct = mystruct.dtype._typeclass.as_ctypes()(data=dat.__array_interface__['data'][0])
+
+    func = init_prog.compile()
+    func(mydat=inp_struct, fill_value=3, N=10, M=5)
+
+    assert np.allclose(dat[0, :], 5) and np.allclose(dat[1, :], 5)
+    assert np.allclose(dat[2, :], 3) and np.allclose(dat[3, :], 3)
+    assert np.allclose(dat[4, :], 4) and np.allclose(dat[5, :], 4)
+    assert np.allclose(dat[6, :], 3) and np.allclose(dat[7, :], 3)
+    assert np.allclose(dat[8, :], 3) and np.allclose(dat[9, :], 3)
+
+
+def test_type_annotation_reference():
+    N = dace.symbol('N')
+
+    @dace.program
+    def ref(A: dace.float64[N], B: dace.float64[N], T: dace.int32, out: dace.float64[N]):
+        ref1: dace.data.ArrayReference(A.dtype, A.shape) = A
+        ref2: dace.data.ArrayReference(A.dtype, A.shape) = B
+        if T <= 0:
+            out[:] = ref1[:] + 1
+        else:
+            out[:] = ref2[:] + 1
+
+    a = np.random.rand(20)
+    a_verif = a.copy()
+    b = np.random.rand(20)
+    b_verif = b.copy()
+    out = np.random.rand(20)
+    out_verif = out.copy()
+
+    ref(a, b, 1, out, N=20)
+    ref.f(a_verif, b_verif, 1, out_verif)
+    assert np.allclose(out, out_verif)
+
+    ref(a, b, -1, out, N=20)
+    ref.f(a_verif, b_verif, -1, out_verif)
+    assert np.allclose(out, out_verif)
+
+
 def test_unset_reference():
     sdfg = dace.SDFG('tester')
     sdfg.add_reference('ref', [20], dace.float64)
@@ -683,6 +747,8 @@ def test_ref2view_reconnection():
 
 
 if __name__ == '__main__':
+    test_frontend_reference()
+    test_type_annotation_reference()
     test_unset_reference()
     test_reference_branch()
     test_reference_sources_pass()