Skip to content

Commit 1dc149d

Browse files
sklamsklam
authored
Merge pull request numba#7359 from luk-f-a/feature/nested_arrays_getitem
Extend support for nested arrays inside numpy records
2 parents d1f15ab + dd1c51e commit 1dc149d

6 files changed

Lines changed: 702 additions & 185 deletions

File tree

‎numba/core/base.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -972,6 +972,11 @@ def sentry_record_alignment(self, rectyp, attr):
972972
if self.strict_alignment:
973973
offset = rectyp.offset(attr)
974974
elemty = rectyp.typeof(attr)
975+
if isinstance(elemty, types.NestedArray):
976+
# For a NestedArray we need to consider the data type of
977+
# elements of the array for alignment, not the array structure
978+
# itself
979+
elemty = elemty.dtype
975980
align = self.get_abi_alignment(self.get_data_type(elemty))
976981
if offset % align:
977982
msg = "{rec}.{attr} of type {type} is not aligned".format(

‎numba/core/datamodel/models.py‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -896,10 +896,6 @@ def __init__(self, dmm, fe_type):
896896
self._be_type = dmm.lookup(fe_type.dtype).get_data_type()
897897
super(NestedArrayModel, self).__init__(dmm, fe_type)
898898

899-
def get_data_type(self):
900-
ret = ir.ArrayType(self._be_type, self._fe_type.nitems)
901-
return ret
902-
903899

904900
@register_default(types.Optional)
905901
class OptionalModel(StructModel):

‎numba/core/typing/arraydecl.py‎

Lines changed: 32 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,7 @@ def generic(self, args, kws):
170170
if out is not None:
171171
return signature(out.result, ary, out.index)
172172

173+
173174
@infer_global(operator.setitem)
174175
class SetItemBuffer(AbstractTemplate):
175176
def generic(self, args, kws):
@@ -185,7 +186,7 @@ def generic(self, args, kws):
185186
return
186187

187188
idx = out.index
188-
res = out.result # res is the result type of the access ary[idx]
189+
res = out.result # res is the result type of the access ary[idx]
189190
if isinstance(res, types.Array):
190191
# Indexing produces an array
191192
if isinstance(val, types.Array):
@@ -220,8 +221,8 @@ def generic(self, args, kws):
220221
else:
221222
return
222223
res = val
223-
elif isinstance(val, types.Array) and val.ndim == 0 \
224-
and self.context.can_convert(val.dtype, res):
224+
elif (isinstance(val, types.Array) and val.ndim == 0
225+
and self.context.can_convert(val.dtype, res)):
225226
# val is an array(T, 0d, O), where T is the type of res, O is order
226227
res = val
227228
else:
@@ -492,7 +493,7 @@ def resolve_take(self, ary, args, kws):
492493
sig = signature(ary.dtype, *args)
493494
elif isinstance(argty, types.Array):
494495
sig = signature(argty.copy(layout='C', dtype=ary.dtype), *args)
495-
elif isinstance(argty, types.List): # 1d lists only
496+
elif isinstance(argty, types.List): # 1d lists only
496497
sig = signature(types.Array(ary.dtype, 1, 'C'), *args)
497498
elif isinstance(argty, types.BaseTuple):
498499
sig = signature(types.Array(ary.dtype, np.ndim(argty), 'C'), *args)
@@ -504,7 +505,15 @@ def generic_resolve(self, ary, attr):
504505
# Resolution of other attributes, for record arrays
505506
if isinstance(ary.dtype, types.Record):
506507
if attr in ary.dtype.fields:
507-
return ary.copy(dtype=ary.dtype.typeof(attr), layout='A')
508+
attr_dtype = ary.dtype.typeof(attr)
509+
if isinstance(attr_dtype, types.NestedArray):
510+
return ary.copy(
511+
dtype=attr_dtype.dtype,
512+
ndim=ary.ndim + attr_dtype.ndim,
513+
layout='A'
514+
)
515+
else:
516+
return ary.copy(dtype=attr_dtype, layout='A')
508517

509518

510519
@infer_getattr
@@ -524,6 +533,7 @@ def resolve_kind(self, ary):
524533
return None # other types not supported yet
525534
return types.StringLiteral(val)
526535

536+
527537
@infer
528538
class StaticGetItemArray(AbstractTemplate):
529539
key = "static_getitem"
@@ -532,10 +542,19 @@ def generic(self, args, kws):
532542
# Resolution of members for record and structured arrays
533543
ary, idx = args
534544
if (isinstance(ary, types.Array) and isinstance(idx, str) and
535-
isinstance(ary.dtype, types.Record)):
545+
isinstance(ary.dtype, types.Record)):
536546
if idx in ary.dtype.fields:
537-
ret = ary.copy(dtype=ary.dtype.typeof(idx), layout='A')
538-
return signature(ret, *args)
547+
attr_dtype = ary.dtype.typeof(idx)
548+
if isinstance(attr_dtype, types.NestedArray):
549+
ret = ary.copy(
550+
dtype=attr_dtype.dtype,
551+
ndim=ary.ndim + attr_dtype.ndim,
552+
layout='A'
553+
)
554+
return signature(ret, *args)
555+
else:
556+
ret = ary.copy(dtype=attr_dtype, layout='A')
557+
return signature(ret, *args)
539558

540559

541560
@infer_getattr
@@ -547,6 +566,7 @@ def generic_resolve(self, record, attr):
547566
assert ret
548567
return ret
549568

569+
550570
@infer
551571
class StaticGetItemRecord(AbstractTemplate):
552572
key = "static_getitem"
@@ -770,6 +790,7 @@ def generic_hetero_real(self, args, kws):
770790
return signature(types.float64, recvr=self.this)
771791
return signature(self.this.dtype, recvr=self.this)
772792

793+
773794
def generic_hetero_always_real(self, args, kws):
774795
assert not args
775796
assert not kws
@@ -779,15 +800,18 @@ def generic_hetero_always_real(self, args, kws):
779800
return signature(self.this.dtype.underlying_float, recvr=self.this)
780801
return signature(self.this.dtype, recvr=self.this)
781802

803+
782804
def generic_index(self, args, kws):
783805
assert not args
784806
assert not kws
785807
return signature(types.intp, recvr=self.this)
786808

809+
787810
def install_array_method(name, generic, prefer_literal=True):
788811
my_attr = {"key": "array." + name, "generic": generic,
789812
"prefer_literal": prefer_literal}
790813
temp_class = type("Array_" + name, (AbstractTemplate,), my_attr)
814+
791815
def array_attribute_attachment(self, ary):
792816
return types.BoundFunction(temp_class, ary)
793817

‎numba/cuda/tests/cudapy/test_record_dtype.py‎

Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,70 @@ def record_read_2d_array(r, a):
9898
('j', np.float32, (3, 2))])
9999

100100

101+
# Functions used for "full array" tests
102+
103+
def record_write_full_array(rec):
104+
rec.j[:, :] = np.ones((3, 2))
105+
106+
107+
def record_write_full_array_alt(rec):
108+
rec['j'][:, :] = np.ones((3, 2))
109+
110+
111+
def recarray_set_record(ary, rec):
112+
ary[0] = rec
113+
114+
115+
def recarray_write_array_of_nestedarray_broadcast(ary):
116+
ary.j[:, :, :] = 1
117+
return ary
118+
119+
120+
def recarray_write_array_of_nestedarray(ary):
121+
ary.j[:, :, :] = np.ones((2, 3, 2))
122+
return ary
123+
124+
125+
def recarray_getitem_return(ary):
126+
return ary[0]
127+
128+
129+
def recarray_getitem_field_return(ary):
130+
return ary['h']
131+
132+
133+
def recarray_getitem_field_return2(ary):
134+
return ary.h
135+
136+
137+
def recarray_getitem_field_return2_2d(ary):
138+
return ary.j
139+
140+
141+
def record_read_array0(ary):
142+
return ary.h[0]
143+
144+
145+
def record_read_array1(ary):
146+
return ary.h[1]
147+
148+
149+
def record_read_whole_array(ary):
150+
return ary.h
151+
152+
153+
def record_read_2d_array00(ary):
154+
return ary.j[0, 0]
155+
156+
157+
def record_read_2d_array10(ary):
158+
return ary.j[1, 0]
159+
160+
161+
def record_read_2d_array01(ary):
162+
return ary.j[0, 1]
163+
164+
101165
class TestRecordDtype(CUDATestCase):
102166

103167
def _createSampleArrays(self):
@@ -279,5 +343,182 @@ def _createSampleArrays(self):
279343
self.samplerec2darr = np.zeros(1, dtype=recordwith2darray)[0]
280344

281345

346+
class TestNestedArrays(CUDATestCase):
347+
348+
# These tests mirror those from
349+
# numba.tests.test_record_dtype.TestNestedArrays added in PR
350+
# #7359: https://github.com/numba/numba/pull/7359
351+
352+
# The code cannot be shared between the two classes without modification,
353+
# as the CUDA test implementations need to be launched (and in some cases
354+
# wrapped in an outer function to handle the return value). Otherwise, the
355+
# code here is kept as similar to that in the equivalent CPU tests as
356+
# possible.
357+
358+
# Reading records / recarrays
359+
360+
def get_cfunc(self, pyfunc, retty):
361+
# Create a host-callable function for testing CUDA device functions
362+
# that get a value from a record array
363+
inner = cuda.jit(device=True)(pyfunc)
364+
365+
@cuda.jit
366+
def outer(arg0, res):
367+
res[0] = inner(arg0)
368+
369+
def host(arg0):
370+
res = np.zeros(1, dtype=retty)
371+
outer[1, 1](arg0, res)
372+
return res[0]
373+
374+
return host
375+
376+
def test_record_read_array(self):
377+
# Test reading from a 1D array within a structured type
378+
nbval = np.recarray(1, dtype=recordwitharray)
379+
nbval[0].h[0] = 15.0
380+
nbval[0].h[1] = 25.0
381+
cfunc = self.get_cfunc(record_read_array0, np.float32)
382+
res = cfunc(nbval[0])
383+
np.testing.assert_equal(res, nbval[0].h[0])
384+
385+
cfunc = self.get_cfunc(record_read_array1, np.float32)
386+
res = cfunc(nbval[0])
387+
np.testing.assert_equal(res, nbval[0].h[1])
388+
389+
def test_record_read_2d_array(self):
390+
# Test reading from a 2D array within a structured type
391+
nbval = np.recarray(1, dtype=recordwith2darray)
392+
nbval[0].j = np.asarray([1.5, 2.5, 3.5, 4.5, 5.5, 6.5],
393+
np.float32).reshape(3, 2)
394+
cfunc = self.get_cfunc(record_read_2d_array00, np.float32)
395+
res = cfunc(nbval[0])
396+
np.testing.assert_equal(res, nbval[0].j[0, 0])
397+
398+
cfunc = self.get_cfunc(record_read_2d_array01, np.float32)
399+
res = cfunc(nbval[0])
400+
np.testing.assert_equal(res, nbval[0].j[0, 1])
401+
402+
cfunc = self.get_cfunc(record_read_2d_array10, np.float32)
403+
res = cfunc(nbval[0])
404+
np.testing.assert_equal(res, nbval[0].j[1, 0])
405+
406+
def test_getitem_idx(self):
407+
# Test __getitem__ with numerical index
408+
409+
# This tests returning a record when passing an array and
410+
# returning the first item when passing a record
411+
nbarr = np.recarray(2, dtype=recordwitharray)
412+
nbarr[0] = np.array([(1, (2, 3))], dtype=recordwitharray)[0]
413+
for arg, retty in [(nbarr, recordwitharray), (nbarr[0], np.int32)]:
414+
pyfunc = recarray_getitem_return
415+
arr_expected = pyfunc(arg)
416+
cfunc = self.get_cfunc(pyfunc, retty)
417+
arr_res = cfunc(arg)
418+
np.testing.assert_equal(arr_res, arr_expected)
419+
420+
# Writing to records / recarrays
421+
422+
@skip_on_cudasim('Structured array attr access not supported in simulator')
423+
def test_set_record(self):
424+
# Test setting an entire record
425+
rec = np.ones(2, dtype=recordwith2darray).view(np.recarray)[0]
426+
nbarr = np.zeros(2, dtype=recordwith2darray).view(np.recarray)
427+
arr = np.zeros(2, dtype=recordwith2darray).view(np.recarray)
428+
pyfunc = recarray_set_record
429+
pyfunc(arr, rec)
430+
kernel = cuda.jit(pyfunc)
431+
kernel[1, 1](nbarr, rec)
432+
np.testing.assert_equal(nbarr, arr)
433+
434+
# Reading and returning arrays from recarrays - the following functions are
435+
# all xfailed because CUDA cannot handle returning arrays from device
436+
# functions (or creating arrays in general).
437+
438+
@unittest.expectedFailure
439+
def test_getitem_idx_2darray(self):
440+
# Test __getitem__ with numerical index
441+
#
442+
# This test returning a record when passing an array and
443+
# return the first item when passing a record
444+
nbarr = np.recarray(2, dtype=recordwith2darray)
445+
nbarr[0] = np.array([(1, ((1,2),(4,5),(2,3)))],
446+
dtype=recordwith2darray)[0]
447+
for arg, retty in [(nbarr, recordwith2darray),
448+
(nbarr[0], (np.float32, (3, 2)))]:
449+
pyfunc = recarray_getitem_field_return2_2d
450+
arr_expected = pyfunc(arg)
451+
cfunc = self.get_cfunc(pyfunc, retty)
452+
arr_res = cfunc(arg)
453+
np.testing.assert_equal(arr_res, arr_expected)
454+
455+
@unittest.expectedFailure
456+
def test_return_getattr_getitem_fieldname(self):
457+
# Test __getitem__ with field name and getattr .field_name
458+
#
459+
# This tests returning a array of nestedarrays when passing an array and
460+
# returning a nestedarray when passing a record
461+
nbarr = np.recarray(2, dtype=recordwitharray)
462+
nbarr[0] = np.array([(1, (2,3))], dtype=recordwitharray)[0]
463+
for arg, retty in [(nbarr, recordwitharray), (nbarr[0], np.float32)]:
464+
for pyfunc in [recarray_getitem_field_return,
465+
recarray_getitem_field_return2]:
466+
arr_expected = pyfunc(arg)
467+
cfunc = self.get_cfunc(pyfunc, retty)
468+
arr_res = cfunc(arg)
469+
np.testing.assert_equal(arr_res, arr_expected)
470+
471+
@unittest.expectedFailure
472+
def test_record_read_arrays(self):
473+
# Test reading from a 1D array within a structured type
474+
nbval = np.recarray(2, dtype=recordwitharray)
475+
nbval[0].h[0] = 15.0
476+
nbval[0].h[1] = 25.0
477+
nbval[1].h[0] = 35.0
478+
nbval[1].h[1] = 45.4
479+
cfunc = self.get_cfunc(record_read_whole_array, np.float32)
480+
res = cfunc(nbval)
481+
np.testing.assert_equal(res, nbval.h)
482+
483+
@unittest.expectedFailure
484+
def test_return_array(self):
485+
# Test getitem record AND array within record and returning it
486+
nbval = np.recarray(2, dtype=recordwitharray)
487+
nbval[0] = np.array([(1, (2,3))], dtype=recordwitharray)[0]
488+
pyfunc = record_read_array0
489+
arr_expected = pyfunc(nbval)
490+
cfunc = self.get_cfunc(pyfunc, np.float32)
491+
arr_res = cfunc(nbval)
492+
np.testing.assert_equal(arr_expected, arr_res)
493+
494+
@skip_on_cudasim('Will unexpectedly pass on cudasim')
495+
@unittest.expectedFailure
496+
def test_set_array(self):
497+
#Test setting an entire array within one record
498+
arr = np.zeros(2, dtype=recordwith2darray).view(np.recarray)
499+
rec = arr[0]
500+
nbarr = np.zeros(2, dtype=recordwith2darray).view(np.recarray)
501+
nbrec = nbarr[0]
502+
for pyfunc in (record_write_full_array, record_write_full_array_alt):
503+
pyfunc(rec)
504+
kernel = cuda.jit(pyfunc)
505+
kernel[1, 1](nbrec)
506+
np.testing.assert_equal(nbarr, arr)
507+
508+
@unittest.expectedFailure
509+
def test_set_arrays(self):
510+
# Test setting an entire array of arrays (multiple records)
511+
arr = np.zeros(2, dtype=recordwith2darray).view(np.recarray)
512+
nbarr = np.zeros(2, dtype=recordwith2darray).view(np.recarray)
513+
for pyfunc in (
514+
recarray_write_array_of_nestedarray_broadcast,
515+
recarray_write_array_of_nestedarray,
516+
):
517+
arr_expected = pyfunc(arr)
518+
cfunc = self.get_cfunc(pyfunc, nbarr.dtype)
519+
arr_res = cfunc(nbarr)
520+
np.testing.assert_equal(arr_res, arr_expected)
521+
522+
282523
if __name__ == '__main__':
283524
unittest.main()

0 commit comments

Comments
 (0)