Merge pull request #7712 from luk-f-a/issue-7693

Fixing issue 7693

numba/cuda/simulator/kernel.py

@@ -148,6 +148,10 @@ def forall(self, ntasks, tpb=0, stream=0, sharedmem=0):
     def overloads(self):
         return FakeOverloadDict()
 
+    @property
+    def py_func(self):
+        return self.fn
+
 
 # Thread emulation
 

numba/cuda/tests/cudapy/test_record_dtype.py

@@ -97,6 +97,10 @@ def record_read_2d_array(r, a):
 recordwith2darray = np.dtype([('i', np.int32),
                               ('j', np.float32, (3, 2))])
 
+nested_array1_dtype = np.dtype([("array1", np.int16, (3,))], align=True)
+
+nested_array2_dtype = np.dtype([("array2", np.int16, (3, 2))], align=True)
+
 
 # Functions used for "full array" tests
 
@@ -117,6 +121,10 @@ def recarray_write_array_of_nestedarray_broadcast(ary):
     return ary
 
 
+def record_setitem_array(rec_source, rec_dest):
+    rec_dest['j'] = rec_source['j']
+
+
 def recarray_write_array_of_nestedarray(ary):
     ary.j[:, :, :] = np.ones((2, 3, 2))
     return ary
@@ -162,6 +170,14 @@ def record_read_2d_array01(ary):
     return ary.j[0, 1]
 
 
+def assign_array_to_nested(dest, src):
+    dest['array1'] = src
+
+
+def assign_array_to_nested_2d(dest, src):
+    dest['array2'] = src
+
+
 class TestRecordDtype(CUDATestCase):
 
     def _createSampleArrays(self):
@@ -403,6 +419,24 @@ def test_record_read_2d_array(self):
         res = cfunc(nbval[0])
         np.testing.assert_equal(res, nbval[0].j[1, 0])
 
+    def test_setitem(self):
+        def gen():
+            nbarr1 = np.recarray(1, dtype=recordwith2darray)
+            nbarr1[0] = np.array([(1, ((1, 2), (4, 5), (2, 3)))],
+                                 dtype=recordwith2darray)[0]
+            nbarr2 = np.recarray(1, dtype=recordwith2darray)
+            nbarr2[0] = np.array([(10, ((10, 20), (40, 50), (20, 30)))],
+                                 dtype=recordwith2darray)[0]
+            return nbarr1[0], nbarr2[0]
+        pyfunc = record_setitem_array
+        pyargs = gen()
+        pyfunc(*pyargs)
+
+        cfunc = cuda.jit(pyfunc)
+        cuargs = gen()
+        cfunc[1, 1](*cuargs)
+        np.testing.assert_equal(pyargs, cuargs)
+
     def test_getitem_idx(self):
         # Test __getitem__ with numerical index
 
@@ -431,6 +465,58 @@ def test_set_record(self):
         kernel[1, 1](nbarr, rec)
         np.testing.assert_equal(nbarr, arr)
 
+    def test_assign_array_to_nested(self):
+        src = (np.arange(3) + 1).astype(np.int16)
+        got = np.zeros(2, dtype=nested_array1_dtype)
+        expected = np.zeros(2, dtype=nested_array1_dtype)
+
+        pyfunc = assign_array_to_nested
+        kernel = cuda.jit(pyfunc)
+
+        kernel[1, 1](got[0], src)
+        pyfunc(expected[0], src)
+
+        np.testing.assert_array_equal(expected, got)
+
+    def test_assign_array_to_nested_2d(self):
+        src = (np.arange(6) + 1).astype(np.int16).reshape((3, 2))
+        got = np.zeros(2, dtype=nested_array2_dtype)
+        expected = np.zeros(2, dtype=nested_array2_dtype)
+
+        pyfunc = assign_array_to_nested_2d
+        kernel = cuda.jit(pyfunc)
+
+        kernel[1, 1](got[0], src)
+        pyfunc(expected[0], src)
+
+        np.testing.assert_array_equal(expected, got)
+
+    def test_issue_7693(self):
+        src_dtype = np.dtype([
+            ("user", np.float64),
+            ("array", np.int16, (3,))],
+            align=True)
+
+        dest_dtype = np.dtype([
+            ("user1", np.float64),
+            ("array1", np.int16, (3,))],
+            align=True)
+
+        @cuda.jit
+        def copy(index, src, dest):
+            dest['user1'] = src[index]['user']
+            dest['array1'] = src[index]['array']
+
+        source = np.zeros(2, dtype=src_dtype)
+        got = np.zeros(2, dtype=dest_dtype)
+        expected = np.zeros(2, dtype=dest_dtype)
+
+        source[0] = (1.2, [1, 2, 3])
+        copy[1, 1](0, source, got[0])
+        copy.py_func(0, source, expected[0])
+
+        np.testing.assert_array_equal(expected, got)
+
     # Reading and returning arrays from recarrays - the following functions are
     # all xfailed because CUDA cannot handle returning arrays from device
     # functions (or creating arrays in general).

numba/np/arrayobj.py

@@ -2548,8 +2548,7 @@ def array_record_getitem(context, builder, sig, args):
 @lower_getattr_generic(types.Record)
 def record_getattr(context, builder, typ, value, attr):
     """
-    Generic getattr() implementation for records: fetch the given
-    record member, i.e. a scalar.
+    Generic getattr() implementation for records: get the given record member.
     """
     context.sentry_record_alignment(typ, attr)
     offset = typ.offset(attr)
@@ -2589,8 +2588,7 @@ def record_getattr(context, builder, typ, value, attr):
 @lower_setattr_generic(types.Record)
 def record_setattr(context, builder, sig, args, attr):
     """
-    Generic setattr() implementation for records: set the given
-    record member, i.e. a scalar.
+    Generic setattr() implementation for records: set the given record member.
     """
     typ, valty = sig.args
     target, val = args
@@ -2599,11 +2597,22 @@ def record_setattr(context, builder, sig, args, attr):
     offset = typ.offset(attr)
     elemty = typ.typeof(attr)
 
-    dptr = cgutils.get_record_member(builder, target, offset,
-                                     context.get_data_type(elemty))
-    val = context.cast(builder, val, valty, elemty)
-    align = None if typ.aligned else 1
-    context.pack_value(builder, elemty, val, dptr, align=align)
+    if isinstance(elemty, types.NestedArray):
+        # Copy the data from the RHS into the nested array
+        val_struct = cgutils.create_struct_proxy(valty)(context, builder,
+                                                        value=args[1])
+        src = val_struct.data
+        dest = cgutils.get_record_member(builder, target, offset,
+                                         src.type.pointee)
+        cgutils.memcpy(builder, dest, src,
+                       context.get_constant(types.intp, elemty.nitems))
+    else:
+        # Set the given scalar record member
+        dptr = cgutils.get_record_member(builder, target, offset,
+                                         context.get_data_type(elemty))
+        val = context.cast(builder, val, valty, elemty)
+        align = None if typ.aligned else 1
+        context.pack_value(builder, elemty, val, dptr, align=align)
 
 
 @lower_builtin('static_getitem', types.Record, types.StringLiteral)

numba/tests/test_record_dtype.py

@@ -240,6 +240,10 @@ def recarray_write_array_of_nestedarray_broadcast(ary):
     return ary
 
 
+def record_setitem_array(rec_source, rec_dest):
+    rec_dest['j'] = rec_source['j']
+
+
 def recarray_write_array_of_nestedarray(ary):
     ary.j[:, :, :] = np.ones((2, 3, 2), dtype=np.float64)
     return ary
@@ -402,6 +406,16 @@ def set_field_slice(arr):
     return arr
 
 
+def assign_array_to_nested(dest):
+    tmp = (np.arange(3) + 1).astype(np.int16)
+    dest['array1'] = tmp
+
+
+def assign_array_to_nested_2d(dest):
+    tmp = (np.arange(6) + 1).astype(np.int16).reshape((3, 2))
+    dest['array2'] = tmp
+
+
 recordtype = np.dtype([('a', np.float64),
                        ('b', np.int16),
                        ('c', np.complex64),
@@ -429,6 +443,10 @@ def set_field_slice(arr):
                               ('p', np.float32, (3, 2, 5, 7)),
                               ('q', 'U10'),])
 
+nested_array1_dtype = np.dtype([("array1", np.int16, (3,))], align=True)
+
+nested_array2_dtype = np.dtype([("array2", np.int16, (3, 2))], align=True)
+
 
 class TestRecordDtypeMakeCStruct(unittest.TestCase):
     def test_two_scalars(self):
@@ -1495,6 +1513,35 @@ def test_set_arrays(self):
             arr_res = cfunc(nbarr)
             np.testing.assert_equal(arr_res, arr_expected)
 
+    def test_setitem(self):
+        def gen():
+            nbarr1 = np.recarray(1, dtype=recordwith2darray)
+            nbarr1[0] = np.array([(1, ((1, 2), (4, 5), (2, 3)))],
+                                 dtype=recordwith2darray)[0]
+            nbarr2 = np.recarray(1, dtype=recordwith2darray)
+            nbarr2[0] = np.array([(10, ((10, 20), (40, 50), (20, 30)))],
+                                 dtype=recordwith2darray)[0]
+            return nbarr1[0], nbarr2[0]
+        pyfunc = record_setitem_array
+        pyargs = gen()
+        pyfunc(*pyargs)
+
+        nbargs = gen()
+        cfunc = self.get_cfunc(pyfunc, tuple((typeof(arg) for arg in nbargs)))
+        cfunc(*nbargs)
+        np.testing.assert_equal(pyargs, nbargs)
+
+    def test_setitem_whole_array_error(self):
+        # Ensure we raise a suitable error when attempting to assign an
+        # array to a whole array's worth of nested arrays.
+        nbarr1 = np.recarray(1, dtype=recordwith2darray)
+        nbarr2 = np.recarray(1, dtype=recordwith2darray)
+        args = (nbarr1, nbarr2)
+        pyfunc = record_setitem_array
+        errmsg = "unsupported array index type"
+        with self.assertRaisesRegex(TypingError, errmsg):
+            self.get_cfunc(pyfunc, tuple((typeof(arg) for arg in args)))
+
     def test_getitem_idx(self):
         # Test __getitem__ with numerical index
 
@@ -1623,6 +1670,52 @@ def test_broadcast_slice(self):
             arr_res = cfunc(arg)
             np.testing.assert_equal(arr_res, arr_expected)
 
+    def test_assign_array_to_nested(self):
+        got = np.zeros(2, dtype=nested_array1_dtype)
+        expected = np.zeros(2, dtype=nested_array1_dtype)
+
+        cfunc = njit(assign_array_to_nested)
+        cfunc(got[0])
+        assign_array_to_nested(expected[0])
+
+        np.testing.assert_array_equal(expected, got)
+
+    def test_assign_array_to_nested_2d(self):
+        got = np.zeros(2, dtype=nested_array2_dtype)
+        expected = np.zeros(2, dtype=nested_array2_dtype)
+
+        cfunc = njit(assign_array_to_nested_2d)
+        cfunc(got[0])
+        assign_array_to_nested_2d(expected[0])
+
+        np.testing.assert_array_equal(expected, got)
+
+    def test_issue_7693(self):
+        src_dtype = np.dtype([
+            ("user", np.float64),
+            ("array", np.int16, (3,))],
+            align=True)
+
+        dest_dtype = np.dtype([
+            ("user1", np.float64),
+            ("array1", np.int16, (3,))],
+            align=True)
+
+        @njit
+        def copy(index, src, dest):
+            dest['user1'] = src[index]['user']
+            dest['array1'] = src[index]['array']
+
+        source = np.zeros(2, dtype=src_dtype)
+        got = np.zeros(2, dtype=dest_dtype)
+        expected = np.zeros(2, dtype=dest_dtype)
+
+        source[0] = (1.2, [1, 2, 3])
+        copy(0, source, got[0])
+        copy.py_func(0, source, expected[0])
+
+        np.testing.assert_array_equal(expected, got)
+
 
 if __name__ == '__main__':
     unittest.main()