From e79939bd8a2deff23194a4cd5d144a9d02b74b8b Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Sat, 5 Oct 2019 10:26:21 -0500 Subject: [PATCH 01/36] Get started on this PR. --- VERSION_INFO | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/VERSION_INFO b/VERSION_INFO index 20f49513e1..0e24a92ffa 100644 --- a/VERSION_INFO +++ b/VERSION_INFO @@ -1 +1 @@ -0.1.11 +0.1.12 From ff8f629dc7263e6d0512c6401012e1d7d65bc475 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Sat, 5 Oct 2019 11:00:24 -0500 Subject: [PATCH 02/36] Bring basic getitems of ListOffsetArray up to speed. --- awkward1/_numba/listoffsetarray.py | 38 +++++++++++++++++++++--------- 1 file changed, 27 insertions(+), 11 deletions(-) diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index cf7fcf736c..18db2c2423 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -30,24 +30,39 @@ def ndim(self): return 1 + self.contenttpe.ndim def getitem_int(self): - return self.getitem_tuple(numba.types.Tuple((numba.int64,))) + return self.contenttpe def getitem_range(self): - return self.getitem_tuple(numba.types.Tuple((numba.types.slice2_type,))) + return self def getitem_tuple(self, wheretpe): - return self.getitem_next(wheretpe, False) + import awkward._numba.listarray + nexttpe = awkward._numba.listarray.ListArrayType(util.index64tpe, util.index64tpe, self, numba.none) + out = nexttpe.getitem_next(wheretpe, False) + return out.getitem_int() def getitem_next(self, wheretpe, isadvanced): if len(wheretpe.types) == 0: return self + headtpe = wheretpe.types[0] + tailtpe = numba.types.Tuple(wheretpe.types[1:]) + if isinstance(headtpe, numba.types.Integer): + return self.contenttpe.getitem_next(tailtpe, isadvanced) + elif isinstance(headtpe, numba.types.SliceType): + return ListOffsetArrayType(self.offsetstpe, self.contenttpe.getitem_next(tailtpe, isadvanced), self.idtpe) + elif isinstance(headtpe, numba.types.EllipsisType): + raise NotImplementedError("ellipsis") + elif isinstance(headtpe, type(numba.typeof(numpy.newaxis))): + raise NotImplementedError("newaxis") + elif isinstance(headtpe, numba.types.Array) and not advanced: + if headtpe.ndim != 1: + raise NotImplementedError("array.ndim != 1") + contenttpe = self.contenttpe.carry().getitem_next(tailtpe, True) + return ListOffsetArrayType(self.startstpe, contenttpe, self.idtpe) + elif isinstance(headtpe, numba.types.Array): + return self.contenttpe.getitem_next(tailtpe, True) else: - headtpe = wheretpe.types[0] - tailtpe = numba.types.Tuple(wheretpe.types[1:]) - if isinstance(headtpe, numba.types.Integer): - return self.contenttpe.getitem_next(tailtpe, isadvanced) - else: - return self.getitem_next(tailtpe, isadvanced) + raise AssertionError(headtpe) def carry(self): import awkward1._numba.listarray @@ -143,11 +158,12 @@ def lower_getitem_int(context, builder, sig, args): val, whereval = args proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) - wherevalp1 = builder.add(whereval, context.get_constant(wheretpe, 1)) if isinstance(wheretpe, numba.types.Literal): - wherevalp1_tpe = wheretpe.literal_type + wherevalp1_tpe = numba.types.IntegerLiteral(wheretpe.literal_value + 1) + wherevalp1 = whereval else: wherevalp1_tpe = wheretpe + wherevalp1 = builder.add(whereval, context.get_constant(wheretpe, 1)) start = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, tpe.offsetstpe.dtype(tpe.offsetstpe, wheretpe), (proxyin.offsets, whereval)) stop = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, tpe.offsetstpe.dtype(tpe.offsetstpe, wherevalp1_tpe), (proxyin.offsets, wherevalp1)) From 1abf837c12f099dccd69b2dd9240226d078fc874 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Sat, 5 Oct 2019 11:10:35 -0500 Subject: [PATCH 03/36] Removed underscores from _numba.util functions. --- awkward1/_numba/content.py | 2 +- awkward1/_numba/listarray.py | 18 +++++------------- awkward1/_numba/listoffsetarray.py | 5 +++++ awkward1/_numba/util.py | 25 +++++++++++++++++++------ tests/test_PR012_listarray_in_numba.py | 2 +- 5 files changed, 31 insertions(+), 21 deletions(-) diff --git a/awkward1/_numba/content.py b/awkward1/_numba/content.py index 47d7b7c2a0..492ba2c031 100644 --- a/awkward1/_numba/content.py +++ b/awkward1/_numba/content.py @@ -36,7 +36,7 @@ def generic(self, args, kwargs): if not isinstance(wheretpe, numba.types.BaseTuple): wheretpe = numba.types.Tuple((wheretpe,)) - wheretpe = util._typing_regularize_slice(wheretpe) + wheretpe = util.typing_regularize_slice(wheretpe) if any(not isinstance(t, (numba.types.Integer, numba.types.SliceType, numba.types.EllipsisType, type(numba.typeof(numpy.newaxis)), numba.types.Array)) for t in wheretpe.types): raise TypeError("only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`), and integer or boolean arrays (possibly jagged) are valid indices") diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index 294174d3d6..dacc4971a6 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -209,18 +209,10 @@ def lower_getitem_range(context, builder, sig, args): @numba.extending.lower_builtin(operator.getitem, ListArrayType, numba.types.BaseTuple) def lower_getitem_tuple(context, builder, sig, args): - rettpe, (arraytpe, wheretpe1) = sig.return_type, sig.args - arrayval, whereval1 = args + rettpe, (arraytpe, wheretpe) = sig.return_type, sig.args + arrayval, whereval = args - wheretpe2 = util._typing_regularize_slice(wheretpe1) - util._regularize_slice.compile(wheretpe2(wheretpe1)) - cres = util._regularize_slice.overloads[(wheretpe1,)] - whereval2 = context.call_internal(builder, cres.fndesc, wheretpe2(wheretpe1), (whereval1,)) - - wheretpe3 = util._typing_broadcast_arrays(wheretpe2) - util.broadcast_arrays.compile(wheretpe3(wheretpe2)) - cres2 = util.broadcast_arrays.overloads[(wheretpe2,)] - whereval3 = context.call_internal(builder, cres2.fndesc, wheretpe3(wheretpe2), (whereval2,)) + wheretpe, whereval = util.preprocess_slicetuple(context, builder, wheretpe, whereval) length = util.arraylen(context, builder, arraytpe, arrayval, totpe=numba.int64) @@ -233,8 +225,8 @@ def lower_getitem_tuple(context, builder, sig, args): proxynext.content = arrayval nextval = proxynext._getvalue() - outtpe = nexttpe.getitem_next(wheretpe3, False) - outval = nexttpe.lower_getitem_next(context, builder, nexttpe, wheretpe3, nextval, whereval3, None) + outtpe = nexttpe.getitem_next(wheretpe, False) + outval = nexttpe.lower_getitem_next(context, builder, nexttpe, wheretpe, nextval, whereval, None) return outtpe.lower_getitem_int(context, builder, rettpe(outtpe, numba.int64), (outval, context.get_constant(numba.int64, 0))) diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index 18db2c2423..cc2638822f 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -198,6 +198,11 @@ def lower_getitem_slice(context, builder, sig, args): context.nrt.incref(builder, rettpe, out) return out +@numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.BaseTuple) +def lower_getitem_tuple(context, builder, sig, args): + raise NotImplementedError + + # def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): # proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) # proxyout = numba.cgutils.create_struct_proxy(arraytpe)(context, builder) diff --git a/awkward1/_numba/util.py b/awkward1/_numba/util.py index b60052bb4a..9b8a37237b 100644 --- a/awkward1/_numba/util.py +++ b/awkward1/_numba/util.py @@ -60,7 +60,7 @@ def newindex64(context, builder, lentpe, lenval): return numba.targets.arrayobj.numpy_empty_nd(context, builder, index64tpe(lentpe), (lenval,)) @numba.jit(nopython=True) -def _shapeat(shapeat, array, at, ndim): +def shapeat(shapeat, array, at, ndim): redat = at - (ndim - array.ndim) if redat < 0: return 1 @@ -97,10 +97,10 @@ def getshape(i, at): if i == -1: return "1" elif isinstance(arrays.types[i], numba.types.Array): - return "_shapeat({}, arrays[{}], {}, {})".format(getshape(i - 1, at), i, at, ndim) + return "shapeat({}, arrays[{}], {}, {})".format(getshape(i - 1, at), i, at, ndim) else: return getshape(i - 1, at) - g = {"_shapeat": _shapeat, "broadcast_to": broadcast_to} + g = {"shapeat": shapeat, "broadcast_to": broadcast_to} exec(""" def impl(arrays): shape = ({}) @@ -109,14 +109,14 @@ def impl(arrays): " ".join("broadcast_to(arrays[{}], shape),".format(at) if isinstance(arrays.types[at], (numba.types.Array, numba.types.Integer)) else "arrays[{}],".format(at) for at in range(len(arrays.types)))), g) return g["impl"] -def _typing_broadcast_arrays(arrays): +def typing_broadcast_arrays(arrays): if not isinstance(arrays, numba.types.BaseTuple) or not any(isinstance(x, numba.types.Array) for x in arrays.types): return arrays else: return numba.types.Tuple([numba.types.Array(numba.int64, 1, "C") if isinstance(t, numba.types.Integer) else t for t in arrays.types]) @numba.generated_jit(nopython=True) -def _regularize_slice(arrays): +def regularize_slice(arrays): if not isinstance(arrays, numba.types.BaseTuple) and isinstance(arrays, (numba.types.ArrayCompatible, numba.types.Sequence)) and isinstance(arrays.dtype, numba.types.Boolean): return lambda arrays: numpy.nonzero(arrays) @@ -141,7 +141,7 @@ def _regularize_slice(arrays): exec(code, g) return g["impl"] -def _typing_regularize_slice(arrays): +def typing_regularize_slice(arrays): out = () if not isinstance(arrays, numba.types.BaseTuple) and isinstance(arrays, (numba.types.ArrayCompatible, numba.types.Sequence)) and isinstance(arrays.dtype, numba.types.Boolean): return numba.types.Tuple(arrays.ndims*(numba.types.Array(numba.int64, 1, "C"),)) @@ -160,3 +160,16 @@ def _typing_regularize_slice(arrays): else: out = out + (t,) return numba.types.Tuple(out) + +def preprocess_slicetuple(context, builder, wheretpe1, whereval1): + wheretpe2 = typing_regularize_slice(wheretpe1) + regularize_slice.compile(wheretpe2(wheretpe1)) + cres = regularize_slice.overloads[(wheretpe1,)] + whereval2 = context.call_internal(builder, cres.fndesc, wheretpe2(wheretpe1), (whereval1,)) + + wheretpe3 = typing_broadcast_arrays(wheretpe2) + broadcast_arrays.compile(wheretpe3(wheretpe2)) + cres2 = broadcast_arrays.overloads[(wheretpe2,)] + whereval3 = context.call_internal(builder, cres2.fndesc, wheretpe3(wheretpe2), (whereval2,)) + + return wheretpe3, whereval3 diff --git a/tests/test_PR012_listarray_in_numba.py b/tests/test_PR012_listarray_in_numba.py index 766ff3d0d0..03ed96b69b 100644 --- a/tests/test_PR012_listarray_in_numba.py +++ b/tests/test_PR012_listarray_in_numba.py @@ -28,7 +28,7 @@ def test_slice_utils(): assert awkward1_numba_util.broadcast_arrays((c, "hello")) == (c, "hello") assert [x.tolist() for x in numpy.broadcast_arrays(a, b)] == [x.tolist() for x in awkward1_numba_util.broadcast_arrays(("hello", a, b))[1:]] - assert [x.tolist() for x in awkward1_numba_util._regularize_slice((numpy.array([0, 1, 2, 3]), numpy.array([[True, False], [False, True], [True, True], [False, False]])))] + assert [x.tolist() for x in awkward1_numba_util.regularize_slice((numpy.array([0, 1, 2, 3]), numpy.array([[True, False], [False, True], [True, True], [False, False]])))] starts = numpy.array([0, 3, 3, 5, 6]) stops = numpy.array([3, 3, 5, 6, 10]) From 7cbc00baf785e7513263bf167ad458e84b5e54bf Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Sat, 5 Oct 2019 12:34:27 -0500 Subject: [PATCH 04/36] Infrastructure for ListOffsetArray.getitem(tuple). --- awkward1/_numba/listarray.py | 12 +------- awkward1/_numba/listoffsetarray.py | 47 +++++++++++++++++++++++++++++- awkward1/_numba/util.py | 25 ++++++++++++---- 3 files changed, 67 insertions(+), 17 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index dacc4971a6..8ce9ba1f1e 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -213,17 +213,7 @@ def lower_getitem_tuple(context, builder, sig, args): arrayval, whereval = args wheretpe, whereval = util.preprocess_slicetuple(context, builder, wheretpe, whereval) - - length = util.arraylen(context, builder, arraytpe, arrayval, totpe=numba.int64) - - nexttpe = ListArrayType(util.index64tpe, util.index64tpe, arraytpe, numba.types.none) - proxynext = numba.cgutils.create_struct_proxy(nexttpe)(context, builder) - proxynext.starts = util.newindex64(context, builder, numba.int64, context.get_constant(numba.int64, 1)) - proxynext.stops = util.newindex64(context, builder, numba.int64, context.get_constant(numba.int64, 1)) - numba.targets.arrayobj.store_item(context, builder, util.index64tpe, context.get_constant(numba.int64, 0), util.arrayptr(context, builder, util.index64tpe, proxynext.starts)) - numba.targets.arrayobj.store_item(context, builder, util.index64tpe, length, util.arrayptr(context, builder, util.index64tpe, proxynext.stops)) - proxynext.content = arrayval - nextval = proxynext._getvalue() + nexttpe, nextval = util.wrap_for_slicetuple(context, builder, arraytpe, arrayval) outtpe = nexttpe.getitem_next(wheretpe, False) outval = nexttpe.lower_getitem_next(context, builder, nexttpe, wheretpe, nextval, whereval, None) diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index cc2638822f..9526f7eb89 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -200,8 +200,53 @@ def lower_getitem_slice(context, builder, sig, args): @numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.BaseTuple) def lower_getitem_tuple(context, builder, sig, args): - raise NotImplementedError + rettpe, (arraytpe, wheretpe) = sig.return_type, sig.args + arrayval, whereval = args + wheretpe, whereval = util.preprocess_slicetuple(context, builder, wheretpe, whereval) + nexttpe, nextval = util.wrap_for_slicetuple(context, builder, arraytpe, arrayval) + + outtpe = nexttpe.getitem_next(wheretpe, False) + outval = nexttpe.lower_getitem_next(context, builder, nexttpe, wheretpe, nextval, whereval, None) + + return outtpe.lower_getitem_int(context, builder, rettpe(outtpe, numba.int64), (outval, conext.get_constant(numba.int64, 0))) + +def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, advanced): + import awkward1._numba.listarray + + if len(wheretpe.types) == 0: + return arrayval + + proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) + lenoffsets = util.arraylen(context, builder, arraytpe.offsetstpe, proxyin.offsets, totpe=numba.int64) + lenstarts = builder.sub(lenoffsets, context.get_constant(numba.int64, 1)) + lencontent = util.arraylen(context, builder, arraytpe.contenttpe, proxyin.content, totpe=numba.int64) + + headtpe = wheretpe.types[0] + tailtpe = numba.types.Tuple(wheretpe.types[1:]) + headval = numba.cgutils.unpack_tuple(builder, whereval)[0] + tailval = context.make_tuple(builder, tailtpe, numba.cgutils.unpack_tuple(builder, whereval)[1:]) + + if isinstance(headtpe, numba.types.Integer): + raise NotImplementedError("ListArray.getitem_next(int)") + + elif isinstance(headtpe, numba.types.SliceType): + raise NotImplementedError("ListArray.getitem_next(slice)") + + elif isinstance(headtpe, numba.types.EllipsisType): + raise NotImplementedError("ListArray.getitem_next(ellipsis)") + + elif isinstance(headtpe, type(numba.typeof(numpy.newaxis))): + raise NotImplementedError("ListArray.getitem_next(newaxis)") + + elif isinstance(headtpe, numba.types.Array) and advanced is None: + raise NotImplementedError("ListArray.getitem_next(Array)") + + elif isinstance(headtpe, numba.types.Array): + raise NotImplementedError("ListArray.getitem_next(advanced Array)") + + else: + raise AssertionError(headtpe) # def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): # proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) diff --git a/awkward1/_numba/util.py b/awkward1/_numba/util.py index 9b8a37237b..e5e7172ce8 100644 --- a/awkward1/_numba/util.py +++ b/awkward1/_numba/util.py @@ -161,11 +161,11 @@ def typing_regularize_slice(arrays): out = out + (t,) return numba.types.Tuple(out) -def preprocess_slicetuple(context, builder, wheretpe1, whereval1): - wheretpe2 = typing_regularize_slice(wheretpe1) - regularize_slice.compile(wheretpe2(wheretpe1)) - cres = regularize_slice.overloads[(wheretpe1,)] - whereval2 = context.call_internal(builder, cres.fndesc, wheretpe2(wheretpe1), (whereval1,)) +def preprocess_slicetuple(context, builder, wheretpe, whereval): + wheretpe2 = typing_regularize_slice(wheretpe) + regularize_slice.compile(wheretpe2(wheretpe)) + cres = regularize_slice.overloads[(wheretpe,)] + whereval2 = context.call_internal(builder, cres.fndesc, wheretpe2(wheretpe), (whereval,)) wheretpe3 = typing_broadcast_arrays(wheretpe2) broadcast_arrays.compile(wheretpe3(wheretpe2)) @@ -173,3 +173,18 @@ def preprocess_slicetuple(context, builder, wheretpe1, whereval1): whereval3 = context.call_internal(builder, cres2.fndesc, wheretpe3(wheretpe2), (whereval2,)) return wheretpe3, whereval3 + +def wrap_for_slicetuple(context, builder, arraytpe, arrayval): + import awkward1._numba.listarray + + length = arraylen(context, builder, arraytpe, arrayval, totpe=numba.int64) + nexttpe = awkward1._numba.listarray.ListArrayType(index64tpe, index64tpe, arraytpe, numba.types.none) + proxynext = numba.cgutils.create_struct_proxy(nexttpe)(context, builder) + proxynext.starts = newindex64(context, builder, numba.int64, context.get_constant(numba.int64, 1)) + proxynext.stops = newindex64(context, builder, numba.int64, context.get_constant(numba.int64, 1)) + numba.targets.arrayobj.store_item(context, builder, index64tpe, context.get_constant(numba.int64, 0), arrayptr(context, builder, index64tpe, proxynext.starts)) + numba.targets.arrayobj.store_item(context, builder, index64tpe, length, arrayptr(context, builder, index64tpe, proxynext.stops)) + proxynext.content = arrayval + nextval = proxynext._getvalue() + + return nexttpe, nextval From 0b5f315bc2244d1f406afcdaba6ea21490a468dc Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Sat, 5 Oct 2019 13:11:30 -0500 Subject: [PATCH 05/36] ListOffsetArray and ListArray's Numba implementations are synchronized. --- awkward1/_numba/listarray.py | 53 ++++++++++++++++++ awkward1/_numba/listoffsetarray.py | 86 +++++++++++++++++++++++++----- 2 files changed, 127 insertions(+), 12 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index 8ce9ba1f1e..b996b2053b 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -301,10 +301,63 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) + proxyout = numba.cgutils.create_struct_proxy(arraytpe)(context, builder) proxyout.starts = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.startstpe(arraytpe.startstpe, carrytpe), (proxyin.starts, carryval)) proxyout.stops = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.stopstpe(arraytpe.stopstpe, carrytpe), (proxyin.stops, carryval)) proxyout.content = proxyin.content + if not isinstance(arraytpe.idtpe, numba.types.NoneType): raise NotImplementedError("array.id is not None") return proxyout._getvalue() + +@numba.typing.templates.infer_getattr +class type_methods(numba.typing.templates.AttributeTemplate): + key = ListArrayType + + def generic_resolve(self, tpe, attr): + if attr == "starts": + return tpe.startstpe + + elif attr == "stops": + return tpe.stopstpe + + elif attr == "content": + return tpe.contenttpe + + elif attr == "id": + if tpe.idtpe == numba.none: + return numba.optional(identity.IdentityType(numba.int32[:, :])) + else: + return tpe.idtpe + +@numba.extending.lower_getattr(ListArrayType, "starts") +def lower_starts(context, builder, tpe, val): + proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) + if context.enable_nrt: + context.nrt.incref(builder, tpe.startstpe, proxyin.starts) + return proxyin.starts + +@numba.extending.lower_getattr(ListArrayType, "stops") +def lower_stops(context, builder, tpe, val): + proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) + if context.enable_nrt: + context.nrt.incref(builder, tpe.stopstpe, proxyin.stops) + return proxyin.stops + +@numba.extending.lower_getattr(ListArrayType, "content") +def lower_content(context, builder, tpe, val): + proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) + if context.enable_nrt: + context.nrt.incref(builder, tpe.contenttpe, proxyin.content) + return proxyin.content + +@numba.extending.lower_getattr(ListArrayType, "id") +def lower_id(context, builder, tpe, val): + proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) + if tpe.idtpe == numba.none: + return context.make_optional_none(builder, identity.IdentityType(numba.int32[:, :])) + else: + if context.enable_nrt: + context.nrt.incref(builder, tpe.idtpe, proxyin.id) + return proxyin.id diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index 9526f7eb89..a7b0395b0b 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -211,6 +211,21 @@ def lower_getitem_tuple(context, builder, sig, args): return outtpe.lower_getitem_int(context, builder, rettpe(outtpe, numba.int64), (outval, conext.get_constant(numba.int64, 0))) +def starts_stops(context, builder, offsetstpe, offsetsval, lenstarts, lenoffsets): + proxyslicestarts = numba.cgutils.create_struct_proxy(numba.types.slice2_type)(context, builder) + proxyslicestarts.start = context.get_constant(numba.intp, 0) + proxyslicestarts.stop = context.get_constant(numba.intp, lenstarts) + proxyslicestarts.step = context.get_constant(numba.intp, 1) + starts = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, offsetstpe(offsetstpe, numba.types.slice2_type), (offsetsval, proxyslicestarts._getvalue())) + + proxyslicestops = numba.cgutils.create_struct_proxy(numba.types.slice2_type)(context, builder) + proxyslicestops.start = context.get_constant(numba.intp, 1) + proxyslicestops.stop = context.get_constant(numba.intp, lenoffsets) + proxyslicestops.step = context.get_constant(numba.intp, 1) + stops = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, offsetstpe(offsetstpe, numba.types.slice2_type), (offsetsval, proxyslicestops._getvalue())) + + return starts, stops + def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, advanced): import awkward1._numba.listarray @@ -222,6 +237,8 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, lenstarts = builder.sub(lenoffsets, context.get_constant(numba.int64, 1)) lencontent = util.arraylen(context, builder, arraytpe.contenttpe, proxyin.content, totpe=numba.int64) + starts, stops = starts_stops(context, builder, tpe.offsetstpe, proxyin.offsets, lenstarts, lenoffsets) + headtpe = wheretpe.types[0] tailtpe = numba.types.Tuple(wheretpe.types[1:]) headval = numba.cgutils.unpack_tuple(builder, whereval)[0] @@ -240,6 +257,46 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, raise NotImplementedError("ListArray.getitem_next(newaxis)") elif isinstance(headtpe, numba.types.Array) and advanced is None: + if headtpe.ndim != 1: + raise NotImplementedError("array.ndim != 1") + + flathead = numba.targets.arrayobj.array_ravel(context, builder, util.int64tep(headtpe), (headval,)) + lenflathead = util.arraylen(context, builder, util.int64tpe, flathead, totpe=numba.int64) + lencarry = builder.mul(lenstarts, lenflathead) + + nextcarry = util.newindex64(context, builder, numba.int64, lencarry) + nextadvanced = util.newindex64(context, builder, numba.int64, lencarry) + nextoffsets = util.newindex64(context, builder, numba.int64, lenoffsets) + + util.call(context, builder, cpu.kernels.awkward_listarray64_getitem_next_array_64, + (util.arrayptr(context, builder, util.index64tpe, nextoffsets), + util.arrayptr(context, builder, util.index64tpe, nextcarry), + util.arrayptr(context, builder, util.index64tpe, nextadvanced), + util.arrayptr(context, builder, arraytpe.offsetstpe, starts), + util.arrayptr(context, builder, arraytpe.offsetstpe, stops), + util.arrayptr(context, builder, util.index64tpe, flathead), + context.get_constant(numba.int64, 0), + context.get_constant(numba.int64, 0), + lenstarts, + lenflathead, + lencontent), + "in {}, indexing error".format(arraytpe.shortname)) + + nexttpe = arraytpe.contenttpe.carry() + nextval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) + + contenttpe = nexttpe.getitem_next(tailtpe, True) + contentval = nexttpe.lower_getitem_next(context, builder, nexttpe, tailtpe, nextval, tailval, nextadvanced) + + if not isinstance(arraytpe.idtpe, numba.types.NoneType): + raise NotImplementedError("array.id is not None") + + outtpe = ListOffsetArrayType(arraytpe.offsetstpe, contenttpe, arraytpe.idtpe) + proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) + proxyout.offsets = nextoffsets + proxyout.content = contentval + return proxyout._getvalue() + raise NotImplementedError("ListArray.getitem_next(Array)") elif isinstance(headtpe, numba.types.Array): @@ -248,18 +305,23 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, else: raise AssertionError(headtpe) -# def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): -# proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) -# proxyout = numba.cgutils.create_struct_proxy(arraytpe)(context, builder) -# proxyout.starts = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.startstpe(arraytpe.startstpe, carrytpe), (proxyin.starts, carryval)) -# proxyout.stops = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.stopstpe(arraytpe.stopstpe, carrytpe), (proxyin.stops, carryval)) -# proxyout.content = proxyin.content -# if not isinstance(arraytpe.idtpe, numba.types.NoneType): -# raise NotImplementedError("array.id is not None") -# outval = proxyout._getvalue() -# if context.enable_nrt: -# context.nrt.incref(builder, arraytpe, outval) -# return outval +def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): + import awkward1._numba.listarray + + proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) + lenoffsets = util.arraylen(context, builder, arraytpe.offsetstpe, proxyin.offsets, totpe=numba.int64) + lenstarts = builder.sub(lenoffsets, context.get_constant(numba.int64, 1)) + + starts, stops = starts_stops(context, builder, tpe.offsetstpe, proxyin.offsets, lenstarts, lenoffsets) + + proxyout = numba.cgutils.create_struct_proxy(awkward1._numba.listarray.ListArray(arraytpe.offsetstpe, arraytpe.offsetstpe, arraytpe.contenttpe, arraytpe.idtpe))(context, builder) + proxyout.starts = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (starts, carryval)) + proxyout.stops = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (stops, carryval)) + proxyout.content = proxyin.content + + if not isinstance(arraytpe.idtpe, numba.types.NoneType): + raise NotImplementedError("array.id is not None") + return proxyout._getvalue() @numba.typing.templates.infer_getattr class type_methods(numba.typing.templates.AttributeTemplate): From 9dcf700db1c3010c09ff81af70e994b5aec61735 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 8 Oct 2019 10:46:12 -0500 Subject: [PATCH 06/36] Pass other getitem types through the tuple-handler. --- awkward1/_numba/listarray.py | 11 +++++ awkward1/_numba/listoffsetarray.py | 11 +++++ .../test_PR015_finish_up_getitem_in_numba.py | 44 +++++++++++++++++++ 3 files changed, 66 insertions(+) create mode 100644 tests/test_PR015_finish_up_getitem_in_numba.py diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index b996b2053b..c1e4665011 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -220,6 +220,17 @@ def lower_getitem_tuple(context, builder, sig, args): return outtpe.lower_getitem_int(context, builder, rettpe(outtpe, numba.int64), (outval, context.get_constant(numba.int64, 0))) +@numba.extending.lower_builtin(operator.getitem, ListArrayType, numba.types.Array) +@numba.extending.lower_builtin(operator.getitem, ListArrayType, numba.types.List) +@numba.extending.lower_builtin(operator.getitem, ListArrayType, numba.types.EllipsisType) +@numba.extending.lower_builtin(operator.getitem, ListArrayType, type(numba.typeof(numpy.newaxis))) +def lower_getitem_other(context, builder, sig, args): + rettpe, (arraytpe, wheretpe) = sig.return_type, sig.args + arrayval, whereval = args + wrappedtpe = numba.types.Tuple((wheretpe,)) + wrappedval = context.make_tuple(builder, wrappedtpe, (whereval,)) + return lower_getitem_tuple(context, builder, rettpe(arraytpe, wrappedtpe), (arrayval, wrappedval)) + def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, advanced): import awkward1._numba.listoffsetarray diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index a7b0395b0b..9d1474efa5 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -211,6 +211,17 @@ def lower_getitem_tuple(context, builder, sig, args): return outtpe.lower_getitem_int(context, builder, rettpe(outtpe, numba.int64), (outval, conext.get_constant(numba.int64, 0))) +@numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.Array) +@numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.List) +@numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.EllipsisType) +@numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, type(numba.typeof(numpy.newaxis))) +def lower_getitem_other(context, builder, sig, args): + rettpe, (arraytpe, wheretpe) = sig.return_type, sig.args + arrayval, whereval = args + wrappedtpe = numba.types.Tuple((wheretpe,)) + wrappedval = context.make_tuple(builder, wrappedtpe, (whereval,)) + return lower_getitem_tuple(context, builder, rettpe(arraytpe, wrappedtpe), (arrayval, wrappedval)) + def starts_stops(context, builder, offsetstpe, offsetsval, lenstarts, lenoffsets): proxyslicestarts = numba.cgutils.create_struct_proxy(numba.types.slice2_type)(context, builder) proxyslicestarts.start = context.get_constant(numba.intp, 0) diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py new file mode 100644 index 0000000000..150a3eb129 --- /dev/null +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -0,0 +1,44 @@ +# BSD 3-Clause License; see https://github.com/jpivarski/awkward-1.0/blob/master/LICENSE + +import sys +import itertools + +import pytest +import numpy +numba = pytest.importorskip("numba") + +import awkward1 + +py27 = (sys.version_info[0] < 3) + +def test_development(): + content = awkward1.layout.NumpyArray(numpy.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) + listarray = awkward1.layout.ListArray64(awkward1.layout.Index64(numpy.array([0, 3, 3, 5, 6])), awkward1.layout.Index64(numpy.array([3, 3, 5, 6, 10])), content) + + @numba.njit + def f1(q): + return q[[2, 0, 0, 1]] + + assert awkward1.tolist(f1(listarray)) == [[3.3, 4.4], [0.0, 1.1, 2.2], [0.0, 1.1, 2.2], []] + + + + + +content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) +offsetsA = numpy.arange(0, 2*3*5 + 5, 5) +offsetsB = numpy.arange(0, 2*3 + 3, 3) +startsA, stopsA = offsetsA[:-1], offsetsA[1:] +startsB, stopsB = offsetsB[:-1], offsetsB[1:] + +listoffsetarrayA64 = awkward1.layout.ListOffsetArray64(awkward1.layout.Index64(offsetsA), content) +listoffsetarrayA32 = awkward1.layout.ListOffsetArray32(awkward1.layout.Index32(offsetsA), content) +listarrayA64 = awkward1.layout.ListArray64(awkward1.layout.Index64(startsA), awkward1.layout.Index64(stopsA), content) +listarrayA32 = awkward1.layout.ListArray32(awkward1.layout.Index32(startsA), awkward1.layout.Index32(stopsA), content) +modelA = numpy.arange(2*3*5*7).reshape(2*3, 5, 7) + +listoffsetarrayB64 = awkward1.layout.ListOffsetArray64(awkward1.layout.Index64(offsetsB), listoffsetarrayA64) +listoffsetarrayB32 = awkward1.layout.ListOffsetArray32(awkward1.layout.Index32(offsetsB), listoffsetarrayA64) +listarrayB64 = awkward1.layout.ListArray64(awkward1.layout.Index64(startsB), awkward1.layout.Index64(stopsB), listarrayA64) +listarrayB32 = awkward1.layout.ListArray32(awkward1.layout.Index32(startsB), awkward1.layout.Index32(stopsB), listarrayA64) +modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) From 1a036001e6110d71ee443f55446e8e111c789d6f Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 8 Oct 2019 11:49:59 -0500 Subject: [PATCH 07/36] A 64-bit ListArray-NumpyArray advanced array index works. --- awkward1/_numba/listarray.py | 93 +++++++++++-------- awkward1/_numba/numpyarray.py | 43 ++++++++- .../test_PR015_finish_up_getitem_in_numba.py | 14 ++- 3 files changed, 109 insertions(+), 41 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index c1e4665011..d89e886047 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -262,50 +262,67 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, elif isinstance(headtpe, type(numba.typeof(numpy.newaxis))): raise NotImplementedError("ListArray.getitem_next(newaxis)") - elif isinstance(headtpe, numba.types.Array) and advanced is None: + elif isinstance(headtpe, numba.types.Array): if headtpe.ndim != 1: raise NotImplementedError("array.ndim != 1") flathead = numba.targets.arrayobj.array_ravel(context, builder, util.index64tpe(headtpe), (headval,)) lenflathead = util.arraylen(context, builder, util.index64tpe, flathead, totpe=numba.int64) - lencarry = builder.mul(lenstarts, lenflathead) - lenoffsets = builder.add(lenstarts, context.get_constant(numba.int64, 1)) - - nextcarry = util.newindex64(context, builder, numba.int64, lencarry) - nextadvanced = util.newindex64(context, builder, numba.int64, lencarry) - nextoffsets = util.newindex64(context, builder, numba.int64, lenoffsets) - - util.call(context, builder, cpu.kernels.awkward_listarray64_getitem_next_array_64, - (util.arrayptr(context, builder, util.index64tpe, nextoffsets), - util.arrayptr(context, builder, util.index64tpe, nextcarry), - util.arrayptr(context, builder, util.index64tpe, nextadvanced), - util.arrayptr(context, builder, arraytpe.startstpe, proxyin.starts), - util.arrayptr(context, builder, arraytpe.stopstpe, proxyin.stops), - util.arrayptr(context, builder, util.index64tpe, flathead), - context.get_constant(numba.int64, 0), - context.get_constant(numba.int64, 0), - lenstarts, - lenflathead, - lencontent), - "in {}, indexing error".format(arraytpe.shortname)) - - nexttpe = arraytpe.contenttpe.carry() - nextval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) - - contenttpe = nexttpe.getitem_next(tailtpe, True) - contentval = nexttpe.lower_getitem_next(context, builder, nexttpe, tailtpe, nextval, tailval, nextadvanced) - - if not isinstance(arraytpe.idtpe, numba.types.NoneType): - raise NotImplementedError("array.id is not None") - - outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(arraytpe.startstpe, contenttpe, arraytpe.idtpe) - proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) - proxyout.offsets = nextoffsets - proxyout.content = contentval - return proxyout._getvalue() - elif isinstance(headtpe, numba.types.Array): - raise NotImplementedError("ListArray.getitem_next(advanced Array)") + if advanced is None: + lencarry = builder.mul(lenstarts, lenflathead) + lenoffsets = builder.add(lenstarts, context.get_constant(numba.int64, 1)) + + nextcarry = util.newindex64(context, builder, numba.int64, lencarry) + nextadvanced = util.newindex64(context, builder, numba.int64, lencarry) + nextoffsets = util.newindex64(context, builder, numba.int64, lenoffsets) + util.call(context, builder, cpu.kernels.awkward_listarray64_getitem_next_array_64, + (util.arrayptr(context, builder, util.index64tpe, nextoffsets), + util.arrayptr(context, builder, util.index64tpe, nextcarry), + util.arrayptr(context, builder, util.index64tpe, nextadvanced), + util.arrayptr(context, builder, arraytpe.startstpe, proxyin.starts), + util.arrayptr(context, builder, arraytpe.stopstpe, proxyin.stops), + util.arrayptr(context, builder, util.index64tpe, flathead), + context.get_constant(numba.int64, 0), + context.get_constant(numba.int64, 0), + lenstarts, + lenflathead, + lencontent), + "in {}, indexing error".format(arraytpe.shortname)) + + nexttpe = arraytpe.contenttpe.carry() + nextval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) + + contenttpe = nexttpe.getitem_next(tailtpe, True) + contentval = nexttpe.lower_getitem_next(context, builder, nexttpe, tailtpe, nextval, tailval, nextadvanced) + + outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(arraytpe.startstpe, contenttpe, arraytpe.idtpe) + proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) + proxyout.offsets = nextoffsets + proxyout.content = contentval + return proxyout._getvalue() + + else: + nextcarry = util.newindex64(context, builder, numba.int64, lenstarts) + nextadvanced = util.newindex64(context, builder, numba.int64, lenstarts) + util.call(context, builder, cpu.kernels.awkward_listarray64_getitem_next_array_advanced_64, + (util.arrayptr(context, builder, util.index64tpe, nextcarry), + util.arrayptr(context, builder, util.index64tpe, nextadvanced), + util.arrayptr(context, builder, arraytpe.startstpe, proxyin.starts), + util.arrayptr(context, builder, arraytpe.stopstpe, proxyin.stops), + util.arrayptr(context, builder, util.index64tpe, flathead), + util.arrayptr(context, builder, util.index64tpe, advanced), + context.get_constant(numba.int64, 0), + context.get_constant(numba.int64, 0), + lenstarts, + lenflathead, + lencontent), + "in {}, indexing error".format(arraytpe.shortname)) + + nexttpe = arraytpe.contenttpe.carry() + nextval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) + + return nexttpe.lower_getitem_next(context, builder, nexttpe, tailtpe, nextval, tailval, nextadvanced) else: raise AssertionError(headtpe) diff --git a/awkward1/_numba/numpyarray.py b/awkward1/_numba/numpyarray.py index 300433a465..625975d1be 100644 --- a/awkward1/_numba/numpyarray.py +++ b/awkward1/_numba/numpyarray.py @@ -115,12 +115,19 @@ def lower_len(context, builder, sig, args): @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.Integer) @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.SliceType) @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.Array) +@numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.List) +@numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.EllipsisType) +@numba.extending.lower_builtin(operator.getitem, NumpyArrayType, type(numba.typeof(numpy.newaxis))) @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.BaseTuple) def lower_getitem(context, builder, sig, args): rettpe, (tpe, wheretpe) = sig.return_type, sig.args val, whereval = args proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) + if isinstance(wheretpe, (numba.types.EllipsisType, type(numba.typeof(numpy.newaxis)))): + wheretpe = numba.types.Tuple((wheretpe,)) + whereval = context.make_tuple(builder, wheretpe, (whereval,)) + if isinstance(rettpe, NumpyArrayType): signature = rettpe.arraytpe(tpe.arraytpe, wheretpe) else: @@ -128,7 +135,7 @@ def lower_getitem(context, builder, sig, args): if isinstance(wheretpe, numba.types.BaseTuple): out = numba.targets.arrayobj.getitem_array_tuple(context, builder, signature, (proxyin.array, whereval)) - elif isinstance(wheretpe, numba.types.Array): + elif isinstance(wheretpe, (numba.types.Array, numba.types.List)): out = numba.targets.arrayobj.fancy_getitem_array(context, builder, signature, (proxyin.array, whereval)) else: out = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, signature, (proxyin.array, whereval)) @@ -140,8 +147,40 @@ def lower_getitem(context, builder, sig, args): else: return out +def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, advanced): + if len(wheretpe.types) == 0: + return arrayval + + headtpe = wheretpe.types[0] + tailtpe = numba.types.Tuple(wheretpe.types[1:]) + headval = numba.cgutils.unpack_tuple(builder, whereval)[0] + tailval = context.make_tuple(builder, tailtpe, numba.cgutils.unpack_tuple(builder, whereval)[1:]) + + if isinstance(headtpe, numba.types.Integer): + raise NotImplementedError("NumpyArray.getitem_next(int)") + + elif isinstance(headtpe, numba.types.SliceType): + raise NotImplementedError("NumpyArray.getitem_next(slice)") + + elif isinstance(headtpe, numba.types.EllipsisType): + raise NotImplementedError("NumpyArray.getitem_next(ellipsis)") + + elif isinstance(headtpe, type(numba.typeof(numpy.newaxis))): + raise NotImplementedError("NumpyArray.getitem_next(newaxis)") + + elif isinstance(headtpe, numba.types.Array): + raise NotImplementedError("NumpyArray.getitem_next(array)") + + else: + raise AssertionError(headtpe) + def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): - return numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe(arraytpe, carrytpe), (arrayval, carryval)) + proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) + proxyout = numba.cgutils.create_struct_proxy(arraytpe)(context, builder) + proxyout.array = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.arraytpe(arraytpe.arraytpe, carrytpe), (proxyin.array, carryval)) + if arraytpe.idtpe != numba.none: + raise NotImplementedError("NumpyArray.id != None") + return proxyout._getvalue() @numba.typing.templates.infer_getattr class type_methods(numba.typing.templates.AttributeTemplate): diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index 150a3eb129..ed34da14a8 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -11,7 +11,7 @@ py27 = (sys.version_info[0] < 3) -def test_development(): +def test_development64(): content = awkward1.layout.NumpyArray(numpy.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) listarray = awkward1.layout.ListArray64(awkward1.layout.Index64(numpy.array([0, 3, 3, 5, 6])), awkward1.layout.Index64(numpy.array([3, 3, 5, 6, 10])), content) @@ -21,9 +21,21 @@ def f1(q): assert awkward1.tolist(f1(listarray)) == [[3.3, 4.4], [0.0, 1.1, 2.2], [0.0, 1.1, 2.2], []] + @numba.njit + def f2(q): + return q[[2, 0, 0, -1], 1] + assert awkward1.tolist(f2(listarray)) == [4.4, 1.1, 1.1, 7.7] +def test_development32(): + content = awkward1.layout.NumpyArray(numpy.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) + listarray = awkward1.layout.ListArray32(awkward1.layout.Index32(numpy.array([0, 3, 3, 5, 6], numpy.int32)), awkward1.layout.Index32(numpy.array([3, 3, 5, 6, 10], numpy.int32)), content) + @numba.njit + def f1(q): + return q[[2, 0, 0, 1]] + + assert awkward1.tolist(f1(listarray)) == [[3.3, 4.4], [0.0, 1.1, 2.2], [0.0, 1.1, 2.2], []] content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) From f1da14f11c7ab45558d4c3fbb8230a54368f60d3 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 8 Oct 2019 14:11:40 -0500 Subject: [PATCH 08/36] Identities that have been tested are correct, but some id assignments are wrong. --- awkward1/_numba/identity.py | 30 ++++++++++++++++ awkward1/_numba/listarray.py | 36 +++++++++++++++---- awkward1/_numba/listoffsetarray.py | 16 ++++++--- awkward1/_numba/numpyarray.py | 13 +++++-- .../test_PR015_finish_up_getitem_in_numba.py | 26 ++++++++++++++ 5 files changed, 107 insertions(+), 14 deletions(-) diff --git a/awkward1/_numba/identity.py b/awkward1/_numba/identity.py index 5bab93f23e..c4030bfb4d 100644 --- a/awkward1/_numba/identity.py +++ b/awkward1/_numba/identity.py @@ -80,6 +80,36 @@ def lower_len(context, builder, sig, args): proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) return numba.targets.arrayobj.array_len(context, builder, numba.types.intp(tpe.arraytpe), (proxyin.array,)) +def lower_range(context, builder, idtpe, wheretpe, idval, whereval): + proxyin = numba.cgutils.create_struct_proxy(idtpe)(context, builder, value=idval) + + proxyout = numba.cgutils.create_struct_proxy(idtpe)(context, builder) + proxyout.ref = proxyin.ref + proxyout.fieldloc = proxyin.fieldloc + proxyout.array = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, idtpe.arraytpe(idtpe.arraytpe, wheretpe), (proxyin.array, whereval)) + outval = proxyout._getvalue() + + if context.enable_nrt: + context.nrt.incref(builder, idtpe.fieldloctpe, proxyout.fieldloc) + context.nrt.incref(builder, idtpe.arraytpe, proxyout.array) + context.nrt.incref(builder, idtpe, outval) + return outval + +def lower_carry(context, builder, idtpe, carrytpe, idval, carryval): + proxyin = numba.cgutils.create_struct_proxy(idtpe)(context, builder, value=idval) + + proxyout = numba.cgutils.create_struct_proxy(idtpe)(context, builder) + proxyout.ref = proxyin.ref + proxyout.fieldloc = proxyin.fieldloc + proxyout.array = numba.targets.arrayobj.fancy_getitem_array(context, builder, idtpe.arraytpe(idtpe.arraytpe, carrytpe), (proxyin.array, carryval)) + outval = proxyout._getvalue() + + if context.enable_nrt: + context.nrt.incref(builder, idtpe.fieldloctpe, proxyout.fieldloc) + context.nrt.incref(builder, idtpe.arraytpe, proxyout.array) + context.nrt.incref(builder, idtpe, outval) + return outval + @numba.typing.templates.infer_getattr class type_methods(numba.typing.templates.AttributeTemplate): key = IdentityType diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index d89e886047..36e8f31d82 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -177,6 +177,8 @@ def lower_getitem_int(context, builder, sig, args): starttpe, stoptpe = numba.int32, numba.int32 elif tpe.bitwidth == 64: starttpe, stoptpe = numba.int64, numba.int64 + else: + raise AssertionError("unrecognized bitwidth") start = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, starttpe(tpe.startstpe, wheretpe), (proxyin.starts, whereval)) stop = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, stoptpe(tpe.stopstpe, wheretpe), (proxyin.stops, whereval)) @@ -185,11 +187,15 @@ def lower_getitem_int(context, builder, sig, args): proxyslice.stop = util.cast(context, builder, tpe, numba.intp, stop) proxyslice.step = context.get_constant(numba.intp, 1) + raise Exception + fcn = context.get_function(operator.getitem, rettpe(tpe.contenttpe, numba.types.slice2_type)) return fcn(builder, (proxyin.content, proxyslice._getvalue())) @numba.extending.lower_builtin(operator.getitem, ListArrayType, numba.types.slice2_type) def lower_getitem_range(context, builder, sig, args): + import awkward1._numba.identity + rettpe, (tpe, wheretpe) = sig.return_type, sig.args val, whereval = args @@ -200,7 +206,7 @@ def lower_getitem_range(context, builder, sig, args): proxyout.stops = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, tpe.stopstpe(tpe.stopstpe, wheretpe), (proxyin.stops, whereval)) proxyout.content = proxyin.content if not isinstance(tpe.idtpe, numba.types.NoneType): - raise NotImplementedError("id is not None") + proxyout.id = awkward1._numba.identity.lower_range(context, builder, tpe.idtpe, wheretpe, proxyin.id, whereval) out = proxyout._getvalue() if context.enable_nrt: @@ -270,13 +276,20 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, lenflathead = util.arraylen(context, builder, util.index64tpe, flathead, totpe=numba.int64) if advanced is None: + if arraytpe.bitwidth == 64: + kernel = cpu.kernels.awkward_listarray64_getitem_next_array_64 + elif arraytpe.bitwidth == 32: + kernel = cpu.kernels.awkward_listarray32_getitem_next_array_64 + else: + raise AssertionError("unrecognized bitwidth") + lencarry = builder.mul(lenstarts, lenflathead) lenoffsets = builder.add(lenstarts, context.get_constant(numba.int64, 1)) nextcarry = util.newindex64(context, builder, numba.int64, lencarry) nextadvanced = util.newindex64(context, builder, numba.int64, lencarry) nextoffsets = util.newindex64(context, builder, numba.int64, lenoffsets) - util.call(context, builder, cpu.kernels.awkward_listarray64_getitem_next_array_64, + util.call(context, builder, kernel, (util.arrayptr(context, builder, util.index64tpe, nextoffsets), util.arrayptr(context, builder, util.index64tpe, nextcarry), util.arrayptr(context, builder, util.index64tpe, nextadvanced), @@ -300,12 +313,21 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) proxyout.offsets = nextoffsets proxyout.content = contentval + if arraytpe.idtpe != numba.none: + proxyout.id = proxyin.id return proxyout._getvalue() else: + if arraytpe.bitwidth == 64: + kernel = cpu.kernels.awkward_listarray64_getitem_next_array_advanced_64 + elif arraytpe.bitwidth == 32: + kernel = cpu.kernels.awkward_listarray32_getitem_next_array_advanced_64 + else: + raise AssertionError("unrecognized bitwidth") + nextcarry = util.newindex64(context, builder, numba.int64, lenstarts) nextadvanced = util.newindex64(context, builder, numba.int64, lenstarts) - util.call(context, builder, cpu.kernels.awkward_listarray64_getitem_next_array_advanced_64, + util.call(context, builder, kernel, (util.arrayptr(context, builder, util.index64tpe, nextcarry), util.arrayptr(context, builder, util.index64tpe, nextadvanced), util.arrayptr(context, builder, arraytpe.startstpe, proxyin.starts), @@ -322,21 +344,23 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, nexttpe = arraytpe.contenttpe.carry() nextval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) + outtpe = nexttpe.getitem_next(tailtpe, True) return nexttpe.lower_getitem_next(context, builder, nexttpe, tailtpe, nextval, tailval, nextadvanced) else: raise AssertionError(headtpe) def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): + import awkward1._numba.identity + proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) proxyout = numba.cgutils.create_struct_proxy(arraytpe)(context, builder) proxyout.starts = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.startstpe(arraytpe.startstpe, carrytpe), (proxyin.starts, carryval)) proxyout.stops = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.stopstpe(arraytpe.stopstpe, carrytpe), (proxyin.stops, carryval)) proxyout.content = proxyin.content - - if not isinstance(arraytpe.idtpe, numba.types.NoneType): - raise NotImplementedError("array.id is not None") + if arraytpe.idtpe != numba.none: + proxyout.id = awkward1._numba.identity.lower_carry(context, builder, arraytpe.idtpe, carrytpe, proxyin.id, carryval) return proxyout._getvalue() @numba.typing.templates.infer_getattr diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index 9d1474efa5..14d1e1851c 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -172,11 +172,11 @@ def lower_getitem_int(context, builder, sig, args): proxyslice.stop = util.cast(context, builder, tpe.offsetstpe.dtype, numba.intp, stop) proxyslice.step = context.get_constant(numba.intp, 1) - fcn = context.get_function(operator.getitem, rettpe(tpe.contenttpe, numba.types.slice2_type)) - return fcn(builder, (proxyin.content, proxyslice._getvalue())) + outtpe = tpe.contenttpe.getitem_range() + return tpe.contenttpe.lower_getitem_range(context, builder, outtpe(tpe.contenttpe, numba.types.slice2_type), (proxyin.content, proxyslice._getvalue())) @numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.slice2_type) -def lower_getitem_slice(context, builder, sig, args): +def lower_getitem_range(context, builder, sig, args): rettpe, (tpe, wheretpe) = sig.return_type, sig.args val, whereval = args @@ -209,6 +209,8 @@ def lower_getitem_tuple(context, builder, sig, args): outtpe = nexttpe.getitem_next(wheretpe, False) outval = nexttpe.lower_getitem_next(context, builder, nexttpe, wheretpe, nextval, whereval, None) + raise Exception + return outtpe.lower_getitem_int(context, builder, rettpe(outtpe, numba.int64), (outval, conext.get_constant(numba.int64, 0))) @numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.Array) @@ -271,6 +273,8 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, if headtpe.ndim != 1: raise NotImplementedError("array.ndim != 1") + raise Exception + flathead = numba.targets.arrayobj.array_ravel(context, builder, util.int64tep(headtpe), (headval,)) lenflathead = util.arraylen(context, builder, util.int64tpe, flathead, totpe=numba.int64) lencarry = builder.mul(lenstarts, lenflathead) @@ -306,10 +310,10 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) proxyout.offsets = nextoffsets proxyout.content = contentval + if outtpe.idtpe != numba.none: + proxyout.id = proxyin.id return proxyout._getvalue() - raise NotImplementedError("ListArray.getitem_next(Array)") - elif isinstance(headtpe, numba.types.Array): raise NotImplementedError("ListArray.getitem_next(advanced Array)") @@ -319,6 +323,8 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): import awkward1._numba.listarray + raise Exception + proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) lenoffsets = util.arraylen(context, builder, arraytpe.offsetstpe, proxyin.offsets, totpe=numba.int64) lenstarts = builder.sub(lenoffsets, context.get_constant(numba.int64, 1)) diff --git a/awkward1/_numba/numpyarray.py b/awkward1/_numba/numpyarray.py index 625975d1be..70b677bdc0 100644 --- a/awkward1/_numba/numpyarray.py +++ b/awkward1/_numba/numpyarray.py @@ -57,7 +57,7 @@ def lower_getitem_int(self): @property def lower_getitem_range(self): - return lower_getitem_range + return lower_getitem @property def lower_getitem_next(self): @@ -143,7 +143,10 @@ def lower_getitem(context, builder, sig, args): if isinstance(rettpe, NumpyArrayType): proxyout = numba.cgutils.create_struct_proxy(rettpe)(context, builder) proxyout.array = out - return proxyout._getvalue() + if rettpe.idtpe != numba.none: + proxyout.id = proxyin.id + outval = proxyout._getvalue() + return outval else: return out @@ -175,11 +178,15 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, raise AssertionError(headtpe) def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): + import awkward1._numba.identity + proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) + proxyout = numba.cgutils.create_struct_proxy(arraytpe)(context, builder) proxyout.array = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.arraytpe(arraytpe.arraytpe, carrytpe), (proxyin.array, carryval)) if arraytpe.idtpe != numba.none: - raise NotImplementedError("NumpyArray.id != None") + proxyout.id = awkward1._numba.identity.lower_carry(context, builder, arraytpe.idtpe, carrytpe, proxyin.id, carryval) + return proxyout._getvalue() @numba.typing.templates.infer_getattr diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index ed34da14a8..c73efd8e85 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -27,6 +27,17 @@ def f2(q): assert awkward1.tolist(f2(listarray)) == [4.4, 1.1, 1.1, 7.7] + @numba.njit + def f3(q): + return q[[2, 0, 0, -1], [-1, -1, 0, 0]] + + assert awkward1.tolist(f3(listarray)) == [4.4, 2.2, 0.0, 6.6] + + listarray.setid() + assert numpy.asarray(f1(listarray).id).tolist() == [[2], [0], [0], [1]] + assert numpy.asarray(f2(listarray).id).tolist() == [[2, 1], [0, 1], [0, 1], [4, 1]] + assert numpy.asarray(f3(listarray).id).tolist() == [[2, 1], [0, 2], [0, 0], [4, 0]] + def test_development32(): content = awkward1.layout.NumpyArray(numpy.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) listarray = awkward1.layout.ListArray32(awkward1.layout.Index32(numpy.array([0, 3, 3, 5, 6], numpy.int32)), awkward1.layout.Index32(numpy.array([3, 3, 5, 6, 10], numpy.int32)), content) @@ -37,6 +48,21 @@ def f1(q): assert awkward1.tolist(f1(listarray)) == [[3.3, 4.4], [0.0, 1.1, 2.2], [0.0, 1.1, 2.2], []] + @numba.njit + def f2(q): + return q[[2, 0, 0, -1], 1] + + assert awkward1.tolist(f2(listarray)) == [4.4, 1.1, 1.1, 7.7] + + @numba.njit + def f3(q): + return q[[2, 0, 0, -1], [-1, -1, 0, 0]] + + assert awkward1.tolist(f3(listarray)) == [4.4, 2.2, 0.0, 6.6] + + listarray.setid() + assert numpy.asarray(f1(listarray).id).tolist() == [[2], [0], [0], [1]] + content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) offsetsB = numpy.arange(0, 2*3 + 3, 3) From 863e16db85ecee805c47fe51723621bc55fd4555 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 8 Oct 2019 14:34:03 -0500 Subject: [PATCH 09/36] Fixing wrong .id assignments. --- awkward1/_numba/identity.py | 31 ++++++++----------------------- awkward1/_numba/listarray.py | 4 ++-- awkward1/_numba/numpyarray.py | 8 +++++--- 3 files changed, 15 insertions(+), 28 deletions(-) diff --git a/awkward1/_numba/identity.py b/awkward1/_numba/identity.py index c4030bfb4d..a2c21bd913 100644 --- a/awkward1/_numba/identity.py +++ b/awkward1/_numba/identity.py @@ -80,35 +80,20 @@ def lower_len(context, builder, sig, args): proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) return numba.targets.arrayobj.array_len(context, builder, numba.types.intp(tpe.arraytpe), (proxyin.array,)) -def lower_range(context, builder, idtpe, wheretpe, idval, whereval): +def lower_getitem_any(context, builder, idtpe, wheretpe, idval, whereval): proxyin = numba.cgutils.create_struct_proxy(idtpe)(context, builder, value=idval) proxyout = numba.cgutils.create_struct_proxy(idtpe)(context, builder) proxyout.ref = proxyin.ref proxyout.fieldloc = proxyin.fieldloc - proxyout.array = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, idtpe.arraytpe(idtpe.arraytpe, wheretpe), (proxyin.array, whereval)) - outval = proxyout._getvalue() - - if context.enable_nrt: - context.nrt.incref(builder, idtpe.fieldloctpe, proxyout.fieldloc) - context.nrt.incref(builder, idtpe.arraytpe, proxyout.array) - context.nrt.incref(builder, idtpe, outval) - return outval - -def lower_carry(context, builder, idtpe, carrytpe, idval, carryval): - proxyin = numba.cgutils.create_struct_proxy(idtpe)(context, builder, value=idval) + if isinstance(wheretpe, numba.types.BaseTuple): + proxyout.array = numba.targets.arrayobj.getitem_array_tuple(context, builder, idtpe.arraytpe(idtpe.arraytpe, wheretpe), (proxyin.array, whereval)) + elif isinstance(wheretpe, numba.types.Array): + proxyout.array = numba.targets.arrayobj.fancy_getitem_array(context, builder, idtpe.arraytpe(idtpe.arraytpe, wheretpe), (proxyin.array, whereval)) + else: + proxyout.array = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, idtpe.arraytpe(idtpe.arraytpe, wheretpe), (proxyin.array, whereval)) - proxyout = numba.cgutils.create_struct_proxy(idtpe)(context, builder) - proxyout.ref = proxyin.ref - proxyout.fieldloc = proxyin.fieldloc - proxyout.array = numba.targets.arrayobj.fancy_getitem_array(context, builder, idtpe.arraytpe(idtpe.arraytpe, carrytpe), (proxyin.array, carryval)) - outval = proxyout._getvalue() - - if context.enable_nrt: - context.nrt.incref(builder, idtpe.fieldloctpe, proxyout.fieldloc) - context.nrt.incref(builder, idtpe.arraytpe, proxyout.array) - context.nrt.incref(builder, idtpe, outval) - return outval + return proxyout._getvalue() @numba.typing.templates.infer_getattr class type_methods(numba.typing.templates.AttributeTemplate): diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index 36e8f31d82..d53dcd9e0b 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -206,7 +206,7 @@ def lower_getitem_range(context, builder, sig, args): proxyout.stops = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, tpe.stopstpe(tpe.stopstpe, wheretpe), (proxyin.stops, whereval)) proxyout.content = proxyin.content if not isinstance(tpe.idtpe, numba.types.NoneType): - proxyout.id = awkward1._numba.identity.lower_range(context, builder, tpe.idtpe, wheretpe, proxyin.id, whereval) + proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, tpe.idtpe, wheretpe, proxyin.id, whereval) out = proxyout._getvalue() if context.enable_nrt: @@ -360,7 +360,7 @@ def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): proxyout.stops = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.stopstpe(arraytpe.stopstpe, carrytpe), (proxyin.stops, carryval)) proxyout.content = proxyin.content if arraytpe.idtpe != numba.none: - proxyout.id = awkward1._numba.identity.lower_carry(context, builder, arraytpe.idtpe, carrytpe, proxyin.id, carryval) + proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, arraytpe.idtpe, carrytpe, proxyin.id, carryval) return proxyout._getvalue() @numba.typing.templates.infer_getattr diff --git a/awkward1/_numba/numpyarray.py b/awkward1/_numba/numpyarray.py index 70b677bdc0..772b57d680 100644 --- a/awkward1/_numba/numpyarray.py +++ b/awkward1/_numba/numpyarray.py @@ -120,6 +120,8 @@ def lower_len(context, builder, sig, args): @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, type(numba.typeof(numpy.newaxis))) @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.BaseTuple) def lower_getitem(context, builder, sig, args): + import awkward1._numba.identity + rettpe, (tpe, wheretpe) = sig.return_type, sig.args val, whereval = args proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) @@ -135,7 +137,7 @@ def lower_getitem(context, builder, sig, args): if isinstance(wheretpe, numba.types.BaseTuple): out = numba.targets.arrayobj.getitem_array_tuple(context, builder, signature, (proxyin.array, whereval)) - elif isinstance(wheretpe, (numba.types.Array, numba.types.List)): + elif isinstance(wheretpe, numba.types.Array): out = numba.targets.arrayobj.fancy_getitem_array(context, builder, signature, (proxyin.array, whereval)) else: out = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, signature, (proxyin.array, whereval)) @@ -144,7 +146,7 @@ def lower_getitem(context, builder, sig, args): proxyout = numba.cgutils.create_struct_proxy(rettpe)(context, builder) proxyout.array = out if rettpe.idtpe != numba.none: - proxyout.id = proxyin.id + proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, rettpe.idtpe, wheretpe, proxyin.id, whereval) outval = proxyout._getvalue() return outval else: @@ -185,7 +187,7 @@ def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): proxyout = numba.cgutils.create_struct_proxy(arraytpe)(context, builder) proxyout.array = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.arraytpe(arraytpe.arraytpe, carrytpe), (proxyin.array, carryval)) if arraytpe.idtpe != numba.none: - proxyout.id = awkward1._numba.identity.lower_carry(context, builder, arraytpe.idtpe, carrytpe, proxyin.id, carryval) + proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, arraytpe.idtpe, carrytpe, proxyin.id, carryval) return proxyout._getvalue() From ba61aa51872811a3d9a30a45809e6d9f1915c48e Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 8 Oct 2019 14:41:26 -0500 Subject: [PATCH 10/36] Now all the .id assignments should be correct. --- awkward1/_numba/listarray.py | 6 +++--- awkward1/_numba/listoffsetarray.py | 12 +++--------- tests/test_PR015_finish_up_getitem_in_numba.py | 2 ++ 3 files changed, 8 insertions(+), 12 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index d53dcd9e0b..7839b8b498 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -205,7 +205,7 @@ def lower_getitem_range(context, builder, sig, args): proxyout.starts = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, tpe.startstpe(tpe.startstpe, wheretpe), (proxyin.starts, whereval)) proxyout.stops = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, tpe.stopstpe(tpe.stopstpe, wheretpe), (proxyin.stops, whereval)) proxyout.content = proxyin.content - if not isinstance(tpe.idtpe, numba.types.NoneType): + if tpe.idtpe != numba.none: proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, tpe.idtpe, wheretpe, proxyin.id, whereval) out = proxyout._getvalue() @@ -313,8 +313,8 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) proxyout.offsets = nextoffsets proxyout.content = contentval - if arraytpe.idtpe != numba.none: - proxyout.id = proxyin.id + if outtpe.idtpe != numba.none: + proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, outtpe.idtpe, util.index64tpe, proxyin.id, flathead) return proxyout._getvalue() else: diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index 14d1e1851c..432ea39281 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -303,15 +303,12 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, contenttpe = nexttpe.getitem_next(tailtpe, True) contentval = nexttpe.lower_getitem_next(context, builder, nexttpe, tailtpe, nextval, tailval, nextadvanced) - if not isinstance(arraytpe.idtpe, numba.types.NoneType): - raise NotImplementedError("array.id is not None") - outtpe = ListOffsetArrayType(arraytpe.offsetstpe, contenttpe, arraytpe.idtpe) proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) proxyout.offsets = nextoffsets proxyout.content = contentval if outtpe.idtpe != numba.none: - proxyout.id = proxyin.id + proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, outtpe.idtpe, util.index64tpe, proxyin.id, flathead) return proxyout._getvalue() elif isinstance(headtpe, numba.types.Array): @@ -323,8 +320,6 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): import awkward1._numba.listarray - raise Exception - proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) lenoffsets = util.arraylen(context, builder, arraytpe.offsetstpe, proxyin.offsets, totpe=numba.int64) lenstarts = builder.sub(lenoffsets, context.get_constant(numba.int64, 1)) @@ -335,9 +330,8 @@ def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): proxyout.starts = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (starts, carryval)) proxyout.stops = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (stops, carryval)) proxyout.content = proxyin.content - - if not isinstance(arraytpe.idtpe, numba.types.NoneType): - raise NotImplementedError("array.id is not None") + if arraytpe.idtpe != numba.none: + proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, arraytpe.idtpe, carrytpe, proxyin.id, carryval) return proxyout._getvalue() @numba.typing.templates.infer_getattr diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index c73efd8e85..c99f058bc0 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -62,6 +62,8 @@ def f3(q): listarray.setid() assert numpy.asarray(f1(listarray).id).tolist() == [[2], [0], [0], [1]] + assert numpy.asarray(f2(listarray).id).tolist() == [[2, 1], [0, 1], [0, 1], [4, 1]] + assert numpy.asarray(f3(listarray).id).tolist() == [[2, 1], [0, 2], [0, 0], [4, 0]] content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) From ffb2b8f42866dae724204d031270ce3767de90c6 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 8 Oct 2019 14:48:45 -0500 Subject: [PATCH 11/36] Cleaned up all TODOs. --- awkward1/_numba/identity.py | 8 ++++++++ awkward1/_numba/listarray.py | 6 ++---- awkward1/_numba/listoffsetarray.py | 4 ---- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/awkward1/_numba/identity.py b/awkward1/_numba/identity.py index a2c21bd913..1a635ef534 100644 --- a/awkward1/_numba/identity.py +++ b/awkward1/_numba/identity.py @@ -83,6 +83,14 @@ def lower_len(context, builder, sig, args): def lower_getitem_any(context, builder, idtpe, wheretpe, idval, whereval): proxyin = numba.cgutils.create_struct_proxy(idtpe)(context, builder, value=idval) + if isinstance(wheretpe, numba.types.Integer): + proxyslice = numba.cgutils.create_struct_proxy(numba.types.slice2_type)(context, builder) + proxyslice.start = util.cast(context, builder, wheretpe, numba.intp, whereval) + proxyslice.stop = builder.add(proxyslice.start, context.get_constant(numba.intp, 1)) + proxyslice.step = context.get_constant(numba.intp, 1) + wheretpe = numba.types.slice2_type + whereval = proxyslice._getvalue() + proxyout = numba.cgutils.create_struct_proxy(idtpe)(context, builder) proxyout.ref = proxyin.ref proxyout.fieldloc = proxyin.fieldloc diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index 7839b8b498..a955832bc5 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -187,10 +187,8 @@ def lower_getitem_int(context, builder, sig, args): proxyslice.stop = util.cast(context, builder, tpe, numba.intp, stop) proxyslice.step = context.get_constant(numba.intp, 1) - raise Exception - - fcn = context.get_function(operator.getitem, rettpe(tpe.contenttpe, numba.types.slice2_type)) - return fcn(builder, (proxyin.content, proxyslice._getvalue())) + outtpe = tpe.contenttpe.getitem_range() + return tpe.contenttpe.lower_getitem_range(context, builder, outtpe(tpe.contenttpe, numba.types.slice2_type), (proxyin.content, proxyslice._getvalue())) @numba.extending.lower_builtin(operator.getitem, ListArrayType, numba.types.slice2_type) def lower_getitem_range(context, builder, sig, args): diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index 432ea39281..751b1a889b 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -209,8 +209,6 @@ def lower_getitem_tuple(context, builder, sig, args): outtpe = nexttpe.getitem_next(wheretpe, False) outval = nexttpe.lower_getitem_next(context, builder, nexttpe, wheretpe, nextval, whereval, None) - raise Exception - return outtpe.lower_getitem_int(context, builder, rettpe(outtpe, numba.int64), (outval, conext.get_constant(numba.int64, 0))) @numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.Array) @@ -273,8 +271,6 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, if headtpe.ndim != 1: raise NotImplementedError("array.ndim != 1") - raise Exception - flathead = numba.targets.arrayobj.array_ravel(context, builder, util.int64tep(headtpe), (headval,)) lenflathead = util.arraylen(context, builder, util.int64tpe, flathead, totpe=numba.int64) lencarry = builder.mul(lenstarts, lenflathead) From 2ec6377409fb2c31be8472dbd42afd501067a87b Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 8 Oct 2019 15:08:40 -0500 Subject: [PATCH 12/36] Transfer to laptop. --- awkward1/_numba/listoffsetarray.py | 3 +++ awkward1/_numba/numpyarray.py | 4 +--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index 751b1a889b..ebbb882724 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -193,6 +193,9 @@ def lower_getitem_range(context, builder, sig, args): proxyout = numba.cgutils.create_struct_proxy(tpe)(context, builder) proxyout.offsets = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, tpe.offsetstpe(tpe.offsetstpe, numba.types.slice2_type), (proxyin.offsets, proxysliceout._getvalue())) proxyout.content = proxyin.content + if tpe.idtpe != numba.none: + proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, tpe.idtpe, wheretpe, proxyin.id, whereval) + out = proxyout._getvalue() if context.enable_nrt: context.nrt.incref(builder, rettpe, out) diff --git a/awkward1/_numba/numpyarray.py b/awkward1/_numba/numpyarray.py index 772b57d680..f7fce2ff87 100644 --- a/awkward1/_numba/numpyarray.py +++ b/awkward1/_numba/numpyarray.py @@ -147,8 +147,7 @@ def lower_getitem(context, builder, sig, args): proxyout.array = out if rettpe.idtpe != numba.none: proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, rettpe.idtpe, wheretpe, proxyin.id, whereval) - outval = proxyout._getvalue() - return outval + return proxyout._getvalue() else: return out @@ -188,7 +187,6 @@ def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): proxyout.array = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.arraytpe(arraytpe.arraytpe, carrytpe), (proxyin.array, carryval)) if arraytpe.idtpe != numba.none: proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, arraytpe.idtpe, carrytpe, proxyin.id, carryval) - return proxyout._getvalue() @numba.typing.templates.infer_getattr From bb64e21e021af204a1f275be99bc114858cececa Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 8 Oct 2019 16:39:42 -0500 Subject: [PATCH 13/36] Use arraydecl to type NumpyArray slices. --- awkward1/_numba/numpyarray.py | 10 +++++++--- tests/test_PR015_finish_up_getitem_in_numba.py | 18 ++++++++++++++++-- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/awkward1/_numba/numpyarray.py b/awkward1/_numba/numpyarray.py index f7fce2ff87..2259ab5e46 100644 --- a/awkward1/_numba/numpyarray.py +++ b/awkward1/_numba/numpyarray.py @@ -25,13 +25,17 @@ def ndim(self): return self.arraytpe.ndim def getitem_int(self): - return self.getitem_tuple(numba.types.Tuple((numba.int64,))) + return self.getitem_tuple(numba.int64) def getitem_range(self): - return self.getitem_tuple(numba.types.Tuple((numba.types.slice2_type,))) + return self.getitem_tuple(numba.types.slice2_type) def getitem_tuple(self, wheretpe): - return self.getitem_next(wheretpe, False) + outtpe = numba.typing.arraydecl.get_array_index_type(self.arraytpe, wheretpe).result + if isinstance(outtpe, numba.types.Array): + return NumpyArrayType(outtpe, self.idtpe) + else: + return outtpe def getitem_next(self, wheretpe, isadvanced): if len(wheretpe.types) > self.arraytpe.ndim: diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index c99f058bc0..2563a57177 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -11,7 +11,7 @@ py27 = (sys.version_info[0] < 3) -def test_development64(): +def test_getitem_array_64(): content = awkward1.layout.NumpyArray(numpy.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) listarray = awkward1.layout.ListArray64(awkward1.layout.Index64(numpy.array([0, 3, 3, 5, 6])), awkward1.layout.Index64(numpy.array([3, 3, 5, 6, 10])), content) @@ -38,7 +38,7 @@ def f3(q): assert numpy.asarray(f2(listarray).id).tolist() == [[2, 1], [0, 1], [0, 1], [4, 1]] assert numpy.asarray(f3(listarray).id).tolist() == [[2, 1], [0, 2], [0, 0], [4, 0]] -def test_development32(): +def test_getitem_array_32(): content = awkward1.layout.NumpyArray(numpy.array([0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9])) listarray = awkward1.layout.ListArray32(awkward1.layout.Index32(numpy.array([0, 3, 3, 5, 6], numpy.int32)), awkward1.layout.Index32(numpy.array([3, 3, 5, 6, 10], numpy.int32)), content) @@ -65,6 +65,20 @@ def f3(q): assert numpy.asarray(f2(listarray).id).tolist() == [[2, 1], [0, 1], [0, 1], [4, 1]] assert numpy.asarray(f3(listarray).id).tolist() == [[2, 1], [0, 2], [0, 0], [4, 0]] +def test_deep_numpy(): + content = awkward1.layout.NumpyArray(numpy.array([[0.0, 1.1], [2.2, 3.3], [4.4, 5.5], [6.6, 7.7], [8.8, 9.9]])) + listarray = awkward1.layout.ListArray64(awkward1.layout.Index64(numpy.array([0, 3, 3])), awkward1.layout.Index64(numpy.array([3, 3, 5])), content) + assert awkward1.tolist(listarray) == [[[0.0, 1.1], [2.2, 3.3], [4.4, 5.5]], [], [[6.6, 7.7], [8.8, 9.9]]] + + # @numba.njit + # def f1(q): + # return q[[2, 0, 0, 1]] + # + # print(f1(content)) + # + # raise Exception + + content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) offsetsB = numpy.arange(0, 2*3 + 3, 3) From c39c0505abf8b515731512472d133be01631c4b9 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 8 Oct 2019 16:48:11 -0500 Subject: [PATCH 14/36] Regularize NumpyArray's slice, just like any other. --- awkward1/_numba/numpyarray.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/awkward1/_numba/numpyarray.py b/awkward1/_numba/numpyarray.py index 2259ab5e46..851cc4b8d1 100644 --- a/awkward1/_numba/numpyarray.py +++ b/awkward1/_numba/numpyarray.py @@ -8,7 +8,7 @@ import numba.typing.ctypes_utils import awkward1.layout -from .._numba import cpu, identity, content +from .._numba import cpu, util, identity, content @numba.extending.typeof_impl.register(awkward1.layout.NumpyArray) def typeof(val, c): @@ -130,9 +130,8 @@ def lower_getitem(context, builder, sig, args): val, whereval = args proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) - if isinstance(wheretpe, (numba.types.EllipsisType, type(numba.typeof(numpy.newaxis)))): - wheretpe = numba.types.Tuple((wheretpe,)) - whereval = context.make_tuple(builder, wheretpe, (whereval,)) + if not isinstance(wheretpe, (numba.types.Integer, numba.types.SliceType)): + wheretpe, whereval = util.preprocess_slicetuple(context, builder, wheretpe, whereval) if isinstance(rettpe, NumpyArrayType): signature = rettpe.arraytpe(tpe.arraytpe, wheretpe) From fd60d1e0c79a3daeb567bd80c5fdfb858f913eff Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 8 Oct 2019 17:27:27 -0500 Subject: [PATCH 15/36] Slicing 2-dimensional NumpyArray. --- awkward1/_numba/listarray.py | 1 + awkward1/_numba/listoffsetarray.py | 1 + awkward1/_numba/numpyarray.py | 6 ++++-- awkward1/_numba/util.py | 20 +++++++++---------- .../test_PR015_finish_up_getitem_in_numba.py | 11 ++++------ 5 files changed, 20 insertions(+), 19 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index a955832bc5..45e2053c7f 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -226,6 +226,7 @@ def lower_getitem_tuple(context, builder, sig, args): @numba.extending.lower_builtin(operator.getitem, ListArrayType, numba.types.Array) @numba.extending.lower_builtin(operator.getitem, ListArrayType, numba.types.List) +@numba.extending.lower_builtin(operator.getitem, ListArrayType, numba.types.ArrayCompatible) @numba.extending.lower_builtin(operator.getitem, ListArrayType, numba.types.EllipsisType) @numba.extending.lower_builtin(operator.getitem, ListArrayType, type(numba.typeof(numpy.newaxis))) def lower_getitem_other(context, builder, sig, args): diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index ebbb882724..378a8b8a07 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -216,6 +216,7 @@ def lower_getitem_tuple(context, builder, sig, args): @numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.Array) @numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.List) +@numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.ArrayCompatible) @numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.EllipsisType) @numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, type(numba.typeof(numpy.newaxis))) def lower_getitem_other(context, builder, sig, args): diff --git a/awkward1/_numba/numpyarray.py b/awkward1/_numba/numpyarray.py index 851cc4b8d1..18697e845b 100644 --- a/awkward1/_numba/numpyarray.py +++ b/awkward1/_numba/numpyarray.py @@ -120,6 +120,7 @@ def lower_len(context, builder, sig, args): @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.SliceType) @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.Array) @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.List) +@numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.ArrayCompatible) @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.EllipsisType) @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, type(numba.typeof(numpy.newaxis))) @numba.extending.lower_builtin(operator.getitem, NumpyArrayType, numba.types.BaseTuple) @@ -131,6 +132,9 @@ def lower_getitem(context, builder, sig, args): proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) if not isinstance(wheretpe, (numba.types.Integer, numba.types.SliceType)): + if not isinstance(wheretpe, numba.types.BaseTuple): + wheretpe = numba.types.Tuple((wheretpe,)) + whereval = context.make_tuple(builder, wheretpe, (whereval,)) wheretpe, whereval = util.preprocess_slicetuple(context, builder, wheretpe, whereval) if isinstance(rettpe, NumpyArrayType): @@ -140,8 +144,6 @@ def lower_getitem(context, builder, sig, args): if isinstance(wheretpe, numba.types.BaseTuple): out = numba.targets.arrayobj.getitem_array_tuple(context, builder, signature, (proxyin.array, whereval)) - elif isinstance(wheretpe, numba.types.Array): - out = numba.targets.arrayobj.fancy_getitem_array(context, builder, signature, (proxyin.array, whereval)) else: out = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, signature, (proxyin.array, whereval)) diff --git a/awkward1/_numba/util.py b/awkward1/_numba/util.py index e5e7172ce8..d37984e3ec 100644 --- a/awkward1/_numba/util.py +++ b/awkward1/_numba/util.py @@ -117,22 +117,22 @@ def typing_broadcast_arrays(arrays): @numba.generated_jit(nopython=True) def regularize_slice(arrays): - if not isinstance(arrays, numba.types.BaseTuple) and isinstance(arrays, (numba.types.ArrayCompatible, numba.types.Sequence)) and isinstance(arrays.dtype, numba.types.Boolean): + if not isinstance(arrays, numba.types.BaseTuple) and isinstance(arrays, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(arrays.dtype, numba.types.Boolean): return lambda arrays: numpy.nonzero(arrays) - elif not isinstance(arrays, numba.types.BaseTuple) or not any(isinstance(t, (numba.types.ArrayCompatible, numba.types.Sequence)) for t in arrays.types): + elif not isinstance(arrays, numba.types.BaseTuple) or not any(isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) for t in arrays.types): return lambda arrays: arrays else: code = "def impl(arrays):\n" indexes = [] for i, t in enumerate(arrays.types): - if isinstance(t, (numba.types.ArrayCompatible, numba.types.Sequence)) and isinstance(t.dtype, numba.types.Boolean): + if isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(t.dtype, numba.types.Boolean): code += " x{} = numpy.nonzero(arrays[{}])\n".format(i, i) indexes.extend(["x{}[{}],".format(i, j) for j in range(arrays.types[i].ndim)]) - elif isinstance(t, (numba.types.ArrayCompatible, numba.types.Sequence)) and isinstance(t.dtype, numba.types.Integer): + elif isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(t.dtype, numba.types.Integer): indexes.append("numpy.asarray(arrays[{}], numpy.int64),".format(i)) - elif isinstance(t, (numba.types.ArrayCompatible, numba.types.Sequence)): + elif isinstance(t, (numba.types.ArrayCompatible, numba.types.List)): raise TypeError("arrays must have boolean or integer type") else: indexes.append("arrays[{}]".format(i)) @@ -143,19 +143,19 @@ def regularize_slice(arrays): def typing_regularize_slice(arrays): out = () - if not isinstance(arrays, numba.types.BaseTuple) and isinstance(arrays, (numba.types.ArrayCompatible, numba.types.Sequence)) and isinstance(arrays.dtype, numba.types.Boolean): + if not isinstance(arrays, numba.types.BaseTuple) and isinstance(arrays, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(arrays.dtype, numba.types.Boolean): return numba.types.Tuple(arrays.ndims*(numba.types.Array(numba.int64, 1, "C"),)) - elif not isinstance(arrays, numba.types.BaseTuple) or not any(isinstance(t, (numba.types.ArrayCompatible, numba.types.Sequence)) for t in arrays.types): + elif not isinstance(arrays, numba.types.BaseTuple) or not any(isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) for t in arrays.types): return arrays else: for t in arrays.types: - if isinstance(t, (numba.types.ArrayCompatible, numba.types.Sequence)) and isinstance(t.dtype, numba.types.Boolean): + if isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(t.dtype, numba.types.Boolean): out = out + t.ndims*(numba.types.Array(numba.int64, 1, "C"),) - elif isinstance(t, (numba.types.ArrayCompatible, numba.types.Sequence)) and isinstance(t.dtype, numba.types.Integer): + elif isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(t.dtype, numba.types.Integer): out = out + (numba.types.Array(numba.int64, 1, "C"),) - elif isinstance(t, (numba.types.ArrayCompatible, numba.types.Sequence)): + elif isinstance(t, (numba.types.ArrayCompatible, numba.types.List)): raise TypeError("arrays must have boolean or integer type") else: out = out + (t,) diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index 2563a57177..dc7ba63315 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -70,14 +70,11 @@ def test_deep_numpy(): listarray = awkward1.layout.ListArray64(awkward1.layout.Index64(numpy.array([0, 3, 3])), awkward1.layout.Index64(numpy.array([3, 3, 5])), content) assert awkward1.tolist(listarray) == [[[0.0, 1.1], [2.2, 3.3], [4.4, 5.5]], [], [[6.6, 7.7], [8.8, 9.9]]] - # @numba.njit - # def f1(q): - # return q[[2, 0, 0, 1]] - # - # print(f1(content)) - # - # raise Exception + @numba.njit + def f1(q): + return q[[2, 0, 0, 1]] + assert awkward1.tolist(f1(content)) == [[4.4, 5.5], [0.0, 1.1], [0.0, 1.1], [2.2, 3.3]] content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) From 6ac17887afdd5bbf30f3214a75c381d4e30079f6 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 8 Oct 2019 17:36:56 -0500 Subject: [PATCH 16/36] test_current() is the platform for developing NumpyArray.getitem_next. --- .../test_PR015_finish_up_getitem_in_numba.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index dc7ba63315..8726babd82 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -76,6 +76,31 @@ def f1(q): assert awkward1.tolist(f1(content)) == [[4.4, 5.5], [0.0, 1.1], [0.0, 1.1], [2.2, 3.3]] + assert awkward1.tolist(f1(listarray)) == [[[6.6, 7.7], [8.8, 9.9]], [[0.0, 1.1], [2.2, 3.3], [4.4, 5.5]], [[0.0, 1.1], [2.2, 3.3], [4.4, 5.5]], []] + + @numba.njit + def f2(q): + return q[[2, 0, 0, -1], 1] + + assert awkward1.tolist(f2(listarray)) == [[8.8, 9.9], [2.2, 3.3], [2.2, 3.3], [8.8, 9.9]] + + @numba.njit + def f3(q): + return q[[2, 0, 0, -1], [1, -1, 0, 0]] + + assert awkward1.tolist(f3(listarray)) == [[8.8, 9.9], [4.4, 5.5], [0.0, 1.1], [6.6, 7.7]] + +def test_current(): + pass + + # @numba.njit + # def f4(q): + # return q[[2, 0, 0, -1], [1, -1, 0, 0], 0] + # + # print(f4(listarray)) + # + # raise Exception + content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) offsetsB = numpy.arange(0, 2*3 + 3, 3) From 1f24f9725a955adaadcdf70518eeeb05f9cc19a6 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 13:07:35 -0500 Subject: [PATCH 17/36] [skip ci] in progress. --- awkward1/_numba/numpyarray.py | 70 ++++++++++++++----- src/libawkward/NumpyArray.cpp | 1 - .../test_PR015_finish_up_getitem_in_numba.py | 21 +++--- 3 files changed, 66 insertions(+), 26 deletions(-) diff --git a/awkward1/_numba/numpyarray.py b/awkward1/_numba/numpyarray.py index 18697e845b..13be1fcb51 100644 --- a/awkward1/_numba/numpyarray.py +++ b/awkward1/_numba/numpyarray.py @@ -39,7 +39,7 @@ def getitem_tuple(self, wheretpe): def getitem_next(self, wheretpe, isadvanced): if len(wheretpe.types) > self.arraytpe.ndim: - raise IndexError("too many indices for array") + raise IndexError("too many dimensions in slice") numreduce = sum(1 if isinstance(x, numba.types.Integer) else 0 for x in wheretpe.types) if numreduce < self.arraytpe.ndim: return NumpyArrayType(numba.types.Array(self.arraytpe.dtype, self.arraytpe.ndim - numreduce, self.arraytpe.layout), self.idtpe) @@ -159,29 +159,67 @@ def lower_getitem(context, builder, sig, args): def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, advanced): if len(wheretpe.types) == 0: return arrayval - headtpe = wheretpe.types[0] - tailtpe = numba.types.Tuple(wheretpe.types[1:]) - headval = numba.cgutils.unpack_tuple(builder, whereval)[0] - tailval = context.make_tuple(builder, tailtpe, numba.cgutils.unpack_tuple(builder, whereval)[1:]) + proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) + + if isinstance(headtpe, numba.types.Array) and advanced is not None: + tailtpe = numba.types.Tuple(wheretpe.types[1:]) + headval = numba.cgutils.unpack_tuple(builder, whereval)[0] + tailval = context.make_tuple(builder, tailtpe, numba.cgutils.unpack_tuple(builder, whereval)[1:]) + if headtpe.ndim != 1: + raise NotImplementedError("array.ndim != 1") + if arraytpe.arraytpe.ndim < 2: + raise TypeError("too many dimensions in slice") + + shapeval = numba.targets.arrayobj.make_array(arraytpe.arraytpe)(context, builder, proxyin.array).shape + lenself, skip = numba.cgutils.unpack_tuple(builder, shapeval)[:2] + lenself = util.cast(context, builder, numba.intp, numba.int64, lenself) + skip = util.cast(context, builder, numba.intp, numba.int64, skip) + + carry = numba.targets.arrayobj.numpy_arange_1(context, builder, util.index64tpe(numba.int64), (lenself,)) + + flathead = numba.targets.arrayobj.array_ravel(context, builder, util.index64tpe(headtpe), (headval,)) + lenflathead = util.arraylen(context, builder, util.index64tpe, flathead, totpe=numba.int64) - if isinstance(headtpe, numba.types.Integer): - raise NotImplementedError("NumpyArray.getitem_next(int)") + nextcarry = util.newindex64(context, builder, numba.int64, lenself) + nextadvanced = util.newindex64(context, builder, numba.int64, lenself) - elif isinstance(headtpe, numba.types.SliceType): - raise NotImplementedError("NumpyArray.getitem_next(slice)") + util.call(context, builder, cpu.kernels.awkward_numpyarray_getitem_next_array_advanced_64, + (util.arrayptr(context, builder, util.index64tpe, nextcarry), + util.arrayptr(context, builder, util.index64tpe, carry), + util.arrayptr(context, builder, util.index64tpe, advanced), + util.arrayptr(context, builder, util.index64tpe, flathead), + lenself, + skip), + "in {}, indexing error".format(arraytpe.shortname)) - elif isinstance(headtpe, numba.types.EllipsisType): - raise NotImplementedError("NumpyArray.getitem_next(ellipsis)") + if len(tailtpe.types) != 0: + raise NotImplementedError("len(tail) != 0") + + + raise Exception("HERE") - elif isinstance(headtpe, type(numba.typeof(numpy.newaxis))): - raise NotImplementedError("NumpyArray.getitem_next(newaxis)") - elif isinstance(headtpe, numba.types.Array): - raise NotImplementedError("NumpyArray.getitem_next(array)") else: - raise AssertionError(headtpe) + proxyslice = numba.cgutils.create_struct_proxy(numba.types.slice2_type)(context, builder) + proxyslice.start = context.get_constant(numba.intp, 0) + proxyslice.stop = util.arraylen(context, builder, arraytpe.arraytpe, proxyin.array, totpe=numba.intp) + proxyslice.step = context.get_constant(numba.intp, 1) + wheretpe = numba.types.Tuple((numba.types.slice2_type,) + wheretpe.types) + whereval = context.make_tuple(builder, wheretpe, [proxyslice._getvalue()] + numba.cgutils.unpack_tuple(builder, whereval)) + + outtpe = numba.typing.arraydecl.get_array_index_type(arraytpe.arraytpe, wheretpe).result + outval = numba.targets.arrayobj.getitem_array_tuple(context, builder, outtpe(arraytpe.arraytpe, wheretpe), (proxyin.array, whereval)) + + if isinstance(outtpe, numba.types.Array): + proxyout = numba.cgutils.create_struct_proxy(NumpyArrayType(outtpe, arraytpe.idtpe))(context, builder) + proxyout.array = outval + if arraytpe.idtpe != numba.none: + proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, arraytpe.idtpe, wheretpe, proxyin.id, whereval) + return proxyout._getvalue() + else: + return out def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): import awkward1._numba.identity diff --git a/src/libawkward/NumpyArray.cpp b/src/libawkward/NumpyArray.cpp index 4a1a9c4906..fa3ca8e679 100644 --- a/src/libawkward/NumpyArray.cpp +++ b/src/libawkward/NumpyArray.cpp @@ -725,7 +725,6 @@ namespace awkward { else { Index64 nextcarry(carry.length()); - Index64 nextadvanced(carry.length()); Error err = awkward_numpyarray_getitem_next_array_advanced_64( nextcarry.ptr().get(), carry.ptr().get(), diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index 8726babd82..e8cba85528 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -91,15 +91,18 @@ def f3(q): assert awkward1.tolist(f3(listarray)) == [[8.8, 9.9], [4.4, 5.5], [0.0, 1.1], [6.6, 7.7]] def test_current(): - pass - - # @numba.njit - # def f4(q): - # return q[[2, 0, 0, -1], [1, -1, 0, 0], 0] - # - # print(f4(listarray)) - # - # raise Exception + content = awkward1.layout.NumpyArray(numpy.array([[0.0, 1.1], [2.2, 3.3], [4.4, 5.5], [6.6, 7.7], [8.8, 9.9]])) + listarray = awkward1.layout.ListArray64(awkward1.layout.Index64(numpy.array([0, 3, 3])), awkward1.layout.Index64(numpy.array([3, 3, 5])), content) + + @numba.njit + def f4(q): + return q[[2, 0, 0, -1], [1, -1, 0, 0], [0, 0, 0, 0]] + + # [[8.8, 9.9], [4.4, 5.5], [0.0, 1.1], [6.6, 7.7]] + + print(awkward1.tolist(f4(listarray))) + + raise Exception content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) From 6db10b86309c0bf7202fdaecba5c9aaad472c8b5 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 15:12:23 -0500 Subject: [PATCH 18/36] Fancy indexes pass from ListArray to NumpyArray in Numba. --- awkward1/_numba/cpu.py | 2 +- awkward1/_numba/numpyarray.py | 58 +-- awkward1/_numba/util.py | 12 + awkward1/signatures/Content_8cpp.xml | 4 +- awkward1/signatures/Identity_8cpp.xml | 10 +- awkward1/signatures/ListOffsetArray_8cpp.xml | 18 +- awkward1/signatures/NumpyArray_8cpp.xml | 6 +- .../signatures/cpu-kernels_2util_8cpp.xml | 69 ++++ awkward1/signatures/getitem_8cpp.xml | 342 +++++++++--------- awkward1/signatures/identity_8cpp.xml | 40 +- awkward1/signatures/libawkward_2util_8cpp.xml | 70 ++++ .../test_PR015_finish_up_getitem_in_numba.py | 20 +- 12 files changed, 425 insertions(+), 226 deletions(-) create mode 100644 awkward1/signatures/cpu-kernels_2util_8cpp.xml create mode 100644 awkward1/signatures/libawkward_2util_8cpp.xml diff --git a/awkward1/_numba/cpu.py b/awkward1/_numba/cpu.py index b56c90e264..660001d5c3 100644 --- a/awkward1/_numba/cpu.py +++ b/awkward1/_numba/cpu.py @@ -58,7 +58,7 @@ def __init__(self, dmm, fe_type): if xfile.find("./compounddef/location").attrib["file"].startswith("src/cpu-kernels"): for xfcn in xfile.findall(".//memberdef[@kind='function']"): name = xfcn.find("./name").text - if hasattr(kernels, name): + if hasattr(kernels, name) and name not in ("success", "failure"): rettype = xfcn.find("./type").text params = [(x.find("./declname").text, x.find("./type").text) for x in xfcn.findall("./param")] getattr(kernels, name).name = name diff --git a/awkward1/_numba/numpyarray.py b/awkward1/_numba/numpyarray.py index 13be1fcb51..5fd86e4435 100644 --- a/awkward1/_numba/numpyarray.py +++ b/awkward1/_numba/numpyarray.py @@ -40,7 +40,10 @@ def getitem_tuple(self, wheretpe): def getitem_next(self, wheretpe, isadvanced): if len(wheretpe.types) > self.arraytpe.ndim: raise IndexError("too many dimensions in slice") - numreduce = sum(1 if isinstance(x, numba.types.Integer) else 0 for x in wheretpe.types) + if isadvanced: + numreduce = sum(1 if isinstance(x, (numba.types.Integer, numba.types.Array)) else 0 for x in wheretpe.types) + else: + numreduce = sum(1 if isinstance(x, numba.types.Integer) else 0 for x in wheretpe.types) if numreduce < self.arraytpe.ndim: return NumpyArrayType(numba.types.Array(self.arraytpe.dtype, self.arraytpe.ndim - numreduce, self.arraytpe.layout), self.idtpe) elif numreduce == self.arraytpe.ndim: @@ -172,34 +175,45 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, raise TypeError("too many dimensions in slice") shapeval = numba.targets.arrayobj.make_array(arraytpe.arraytpe)(context, builder, proxyin.array).shape - lenself, skip = numba.cgutils.unpack_tuple(builder, shapeval)[:2] - lenself = util.cast(context, builder, numba.intp, numba.int64, lenself) - skip = util.cast(context, builder, numba.intp, numba.int64, skip) + shapeunpacked = numba.cgutils.unpack_tuple(builder, shapeval) + lenself, skip = shapeunpacked[:2] + lennext = builder.mul(lenself, skip) carry = numba.targets.arrayobj.numpy_arange_1(context, builder, util.index64tpe(numba.int64), (lenself,)) - flathead = numba.targets.arrayobj.array_ravel(context, builder, util.index64tpe(headtpe), (headval,)) - lenflathead = util.arraylen(context, builder, util.index64tpe, flathead, totpe=numba.int64) + flathead = numba.targets.arrayobj.array_flatten(context, builder, util.index64tpe(headtpe), (headval,)) + lenflathead = util.arraylen(context, builder, util.index64tpe, flathead) - nextcarry = util.newindex64(context, builder, numba.int64, lenself) - nextadvanced = util.newindex64(context, builder, numba.int64, lenself) + util.call(context, builder, cpu.kernels.awkward_regularize_arrayslice_64, + (util.arrayptr(context, builder, util.index64tpe, flathead), + util.cast(context, builder, numba.intp, numba.int64, lenflathead), + util.cast(context, builder, numba.intp, numba.int64, skip)), + "in {}, indexing error".format(arraytpe.shortname)) + nextcarry = util.newindex64(context, builder, numba.intp, lenself) util.call(context, builder, cpu.kernels.awkward_numpyarray_getitem_next_array_advanced_64, (util.arrayptr(context, builder, util.index64tpe, nextcarry), util.arrayptr(context, builder, util.index64tpe, carry), util.arrayptr(context, builder, util.index64tpe, advanced), util.arrayptr(context, builder, util.index64tpe, flathead), - lenself, - skip), + util.cast(context, builder, numba.intp, numba.int64, lenself), + util.cast(context, builder, numba.intp, numba.int64, skip)), "in {}, indexing error".format(arraytpe.shortname)) - if len(tailtpe.types) != 0: - raise NotImplementedError("len(tail) != 0") - + nextshapetpe = numba.types.UniTuple(numba.intp, arraytpe.arraytpe.ndim - 1) + nextshapeval = context.make_tuple(builder, nextshapetpe, [lennext] + shapeunpacked[2:]) - raise Exception("HERE") + nextarraytpe = numba.types.Array(arraytpe.arraytpe.dtype, arraytpe.arraytpe.ndim - 1, arraytpe.arraytpe.layout) + nextarrayval = numba.targets.arrayobj.array_reshape(context, builder, nextarraytpe(arraytpe.arraytpe, nextshapetpe), (proxyin.array, nextshapeval)) + nexttpe = NumpyArrayType(nextarraytpe, arraytpe.idtpe) + proxynext = numba.cgutils.create_struct_proxy(nexttpe)(context, builder) + proxynext.array = nextarrayval + if arraytpe.idtpe != numba.none: + proxynext.id = proxyin.id + outval = lower_carry(context, builder, nexttpe, util.index64tpe, proxynext._getvalue(), nextcarry) + return lower_getitem_next(context, builder, nexttpe, tailtpe, outval, tailval, advanced) else: proxyslice = numba.cgutils.create_struct_proxy(numba.types.slice2_type)(context, builder) @@ -212,14 +226,14 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, outtpe = numba.typing.arraydecl.get_array_index_type(arraytpe.arraytpe, wheretpe).result outval = numba.targets.arrayobj.getitem_array_tuple(context, builder, outtpe(arraytpe.arraytpe, wheretpe), (proxyin.array, whereval)) - if isinstance(outtpe, numba.types.Array): - proxyout = numba.cgutils.create_struct_proxy(NumpyArrayType(outtpe, arraytpe.idtpe))(context, builder) - proxyout.array = outval - if arraytpe.idtpe != numba.none: - proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, arraytpe.idtpe, wheretpe, proxyin.id, whereval) - return proxyout._getvalue() - else: - return out + if isinstance(outtpe, numba.types.Array): + proxyout = numba.cgutils.create_struct_proxy(NumpyArrayType(outtpe, arraytpe.idtpe))(context, builder) + proxyout.array = outval + if arraytpe.idtpe != numba.none: + proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, arraytpe.idtpe, wheretpe, proxyin.id, whereval) + return proxyout._getvalue() + else: + return out def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): import awkward1._numba.identity diff --git a/awkward1/_numba/util.py b/awkward1/_numba/util.py index d37984e3ec..6004f1b23b 100644 --- a/awkward1/_numba/util.py +++ b/awkward1/_numba/util.py @@ -1,15 +1,27 @@ # BSD 3-Clause License; see https://github.com/jpivarski/awkward-1.0/blob/master/LICENSE +import sys + import numpy import numba import llvmlite.ir.types from .._numba import cpu +py27 = (sys.version_info[0] < 3) + RefType = numba.int64 index64tpe = numba.types.Array(numba.int64, 1, "C") +if not py27: + exec(""" +def debug(context, builder, *args): + assert len(args) % 2 == 0 + tpes, vals = args[0::2], args[1::2] + context.get_function(print, numba.none(*tpes))(builder, tuple(vals)) +""", globals()) + def cast(context, builder, fromtpe, totpe, val): if fromtpe.bitwidth < totpe.bitwidth: return builder.zext(val, context.get_value_type(totpe)) diff --git a/awkward1/signatures/Content_8cpp.xml b/awkward1/signatures/Content_8cpp.xml index 20ca96ade8..6d83197627 100644 --- a/awkward1/signatures/Content_8cpp.xml +++ b/awkward1/signatures/Content_8cpp.xml @@ -2,7 +2,7 @@ Content.cpp - awkward/ListOffsetArray.h + awkward/ListArray.h awkward/Content.h @@ -17,7 +17,7 @@ - + awkward diff --git a/awkward1/signatures/Identity_8cpp.xml b/awkward1/signatures/Identity_8cpp.xml index f84f47aae1..c4e354d99f 100644 --- a/awkward1/signatures/Identity_8cpp.xml +++ b/awkward1/signatures/Identity_8cpp.xml @@ -9,6 +9,7 @@ type_traits awkward/cpu-kernels/identity.h awkward/cpu-kernels/getitem.h + awkward/Slice.h awkward/Identity.h @@ -45,13 +46,18 @@ + + - + + + + @@ -70,7 +76,7 @@ - + diff --git a/awkward1/signatures/ListOffsetArray_8cpp.xml b/awkward1/signatures/ListOffsetArray_8cpp.xml index 664fd9e9e5..b364c480ee 100644 --- a/awkward1/signatures/ListOffsetArray_8cpp.xml +++ b/awkward1/signatures/ListOffsetArray_8cpp.xml @@ -5,9 +5,14 @@ sstream type_traits awkward/cpu-kernels/identity.h + awkward/cpu-kernels/getitem.h + awkward/Slice.h awkward/ListArray.h awkward/ListOffsetArray.h + + + @@ -15,12 +20,15 @@ + + + - + @@ -36,6 +44,10 @@ + + + + awkward @@ -60,7 +72,7 @@ - + @@ -82,7 +94,7 @@ - + diff --git a/awkward1/signatures/NumpyArray_8cpp.xml b/awkward1/signatures/NumpyArray_8cpp.xml index 17cd6789b8..91c7816f51 100644 --- a/awkward1/signatures/NumpyArray_8cpp.xml +++ b/awkward1/signatures/NumpyArray_8cpp.xml @@ -74,7 +74,7 @@ - + const std::vector< ssize_t > @@ -91,7 +91,7 @@ - + const std::vector< ssize_t > @@ -108,7 +108,7 @@ - + diff --git a/awkward1/signatures/cpu-kernels_2util_8cpp.xml b/awkward1/signatures/cpu-kernels_2util_8cpp.xml new file mode 100644 index 0000000000..6ed44e0983 --- /dev/null +++ b/awkward1/signatures/cpu-kernels_2util_8cpp.xml @@ -0,0 +1,69 @@ + + + + util.cpp + cstring + awkward/cpu-kernels/util.h + + + + + + + + + + + + + + + + + + + Error + Error success + () + success + + + + + + + + + + Error + Error failure + (const char *str, int64_t location, int64_t attempt) + failure + + const char * + str + + + int64_t + location + + + int64_t + attempt + + + + + + + + + + + + + + + + + diff --git a/awkward1/signatures/getitem_8cpp.xml b/awkward1/signatures/getitem_8cpp.xml index 9374d3800f..4a89a4d236 100644 --- a/awkward1/signatures/getitem_8cpp.xml +++ b/awkward1/signatures/getitem_8cpp.xml @@ -86,7 +86,7 @@ - + Error @@ -111,16 +111,16 @@ - + - + typename T - void - void awkward_slicearray_ravel + Error + Error awkward_slicearray_ravel (T *toptr, const T *fromptr, int64_t ndim, const int64_t *shape, const int64_t *strides) awkward_slicearray_ravel @@ -149,11 +149,11 @@ - + - - void - void awkward_slicearray_ravel_64 + + Error + Error awkward_slicearray_ravel_64 (int64_t *toptr, const int64_t *fromptr, int64_t ndim, const int64_t *shape, const int64_t *strides) awkward_slicearray_ravel_64 @@ -182,16 +182,16 @@ - + - + typename T - void - void awkward_carry_arange + Error + Error awkward_carry_arange (T *toptr, int64_t length) awkward_carry_arange @@ -208,11 +208,11 @@ - + - - void - void awkward_carry_arange_64 + + Error + Error awkward_carry_arange_64 (int64_t *toptr, int64_t length) awkward_carry_arange_64 @@ -229,7 +229,7 @@ - + @@ -278,7 +278,7 @@ - + Error @@ -319,7 +319,7 @@ - + Error @@ -360,16 +360,16 @@ - + - + typename T - void - void awkward_numpyarray_contiguous_init + Error + Error awkward_numpyarray_contiguous_init (T *toptr, int64_t skip, int64_t stride) awkward_numpyarray_contiguous_init @@ -390,11 +390,11 @@ - + - - void - void awkward_numpyarray_contiguous_init_64 + + Error + Error awkward_numpyarray_contiguous_init_64 (int64_t *toptr, int64_t skip, int64_t stride) awkward_numpyarray_contiguous_init_64 @@ -415,16 +415,16 @@ - + - + typename T - void - void awkward_numpyarray_contiguous_copy + Error + Error awkward_numpyarray_contiguous_copy (uint8_t *toptr, const uint8_t *fromptr, int64_t len, int64_t stride, int64_t offset, const T *pos) awkward_numpyarray_contiguous_copy @@ -457,11 +457,11 @@ - + - - void - void awkward_numpyarray_contiguous_copy_64 + + Error + Error awkward_numpyarray_contiguous_copy_64 (uint8_t *toptr, const uint8_t *fromptr, int64_t len, int64_t stride, int64_t offset, const int64_t *pos) awkward_numpyarray_contiguous_copy_64 @@ -494,16 +494,16 @@ - + - + typename T - void - void awkward_numpyarray_contiguous_next + Error + Error awkward_numpyarray_contiguous_next (T *topos, const T *frompos, int64_t len, int64_t skip, int64_t stride) awkward_numpyarray_contiguous_next @@ -532,11 +532,11 @@ - + - - void - void awkward_numpyarray_contiguous_next_64 + + Error + Error awkward_numpyarray_contiguous_next_64 (int64_t *topos, const int64_t *frompos, int64_t len, int64_t skip, int64_t stride) awkward_numpyarray_contiguous_next_64 @@ -565,16 +565,16 @@ - + - + typename T - void - void awkward_numpyarray_getitem_next_null + Error + Error awkward_numpyarray_getitem_next_null (uint8_t *toptr, const uint8_t *fromptr, int64_t len, int64_t stride, int64_t offset, const T *pos) awkward_numpyarray_getitem_next_null @@ -607,11 +607,11 @@ - + - - void - void awkward_numpyarray_getitem_next_null_64 + + Error + Error awkward_numpyarray_getitem_next_null_64 (uint8_t *toptr, const uint8_t *fromptr, int64_t len, int64_t stride, int64_t offset, const int64_t *pos) awkward_numpyarray_getitem_next_null_64 @@ -644,16 +644,16 @@ - + - + typename T - void - void awkward_numpyarray_getitem_next_at + Error + Error awkward_numpyarray_getitem_next_at (T *nextcarryptr, const T *carryptr, int64_t lencarry, int64_t skip, int64_t at) awkward_numpyarray_getitem_next_at @@ -682,11 +682,11 @@ - + - - void - void awkward_numpyarray_getitem_next_at_64 + + Error + Error awkward_numpyarray_getitem_next_at_64 (int64_t *nextcarryptr, const int64_t *carryptr, int64_t lencarry, int64_t skip, int64_t at) awkward_numpyarray_getitem_next_at_64 @@ -715,16 +715,16 @@ - + - + typename T - void - void awkward_numpyarray_getitem_next_range + Error + Error awkward_numpyarray_getitem_next_range (T *nextcarryptr, const T *carryptr, int64_t lencarry, int64_t lenhead, int64_t skip, int64_t start, int64_t step) awkward_numpyarray_getitem_next_range @@ -761,11 +761,11 @@ - + - - void - void awkward_numpyarray_getitem_next_range_64 + + Error + Error awkward_numpyarray_getitem_next_range_64 (int64_t *nextcarryptr, const int64_t *carryptr, int64_t lencarry, int64_t lenhead, int64_t skip, int64_t start, int64_t step) awkward_numpyarray_getitem_next_range_64 @@ -802,16 +802,16 @@ - + - + typename T - void - void awkward_numpyarray_getitem_next_range_advanced + Error + Error awkward_numpyarray_getitem_next_range_advanced (T *nextcarryptr, T *nextadvancedptr, const T *carryptr, const T *advancedptr, int64_t lencarry, int64_t lenhead, int64_t skip, int64_t start, int64_t step) awkward_numpyarray_getitem_next_range_advanced @@ -856,11 +856,11 @@ - + - - void - void awkward_numpyarray_getitem_next_range_advanced_64 + + Error + Error awkward_numpyarray_getitem_next_range_advanced_64 (int64_t *nextcarryptr, int64_t *nextadvancedptr, const int64_t *carryptr, const int64_t *advancedptr, int64_t lencarry, int64_t lenhead, int64_t skip, int64_t start, int64_t step) awkward_numpyarray_getitem_next_range_advanced_64 @@ -905,16 +905,16 @@ - + - + typename T - void - void awkward_numpyarray_getitem_next_array + Error + Error awkward_numpyarray_getitem_next_array (T *nextcarryptr, T *nextadvancedptr, const T *carryptr, const T *flatheadptr, int64_t lencarry, int64_t lenflathead, int64_t skip) awkward_numpyarray_getitem_next_array @@ -951,11 +951,11 @@ - + - - void - void awkward_numpyarray_getitem_next_array_64 + + Error + Error awkward_numpyarray_getitem_next_array_64 (int64_t *nextcarryptr, int64_t *nextadvancedptr, const int64_t *carryptr, const int64_t *flatheadptr, int64_t lencarry, int64_t lenflathead, int64_t skip) awkward_numpyarray_getitem_next_array_64 @@ -992,16 +992,16 @@ - + - + typename T - void - void awkward_numpyarray_getitem_next_array_advanced + Error + Error awkward_numpyarray_getitem_next_array_advanced (T *nextcarryptr, const T *carryptr, const T *advancedptr, const T *flatheadptr, int64_t lencarry, int64_t skip) awkward_numpyarray_getitem_next_array_advanced @@ -1034,11 +1034,11 @@ - + - - void - void awkward_numpyarray_getitem_next_array_advanced_64 + + Error + Error awkward_numpyarray_getitem_next_array_advanced_64 (int64_t *nextcarryptr, const int64_t *carryptr, const int64_t *advancedptr, const int64_t *flatheadptr, int64_t lencarry, int64_t skip) awkward_numpyarray_getitem_next_array_advanced_64 @@ -1071,7 +1071,7 @@ - + @@ -1120,7 +1120,7 @@ - + Error @@ -1161,7 +1161,7 @@ - + Error @@ -1202,16 +1202,16 @@ - + - + typename C - void - void awkward_listarray_getitem_next_range_carrylength + Error + Error awkward_listarray_getitem_next_range_carrylength (int64_t *carrylength, const C *fromstarts, const C *fromstops, int64_t lenstarts, int64_t startsoffset, int64_t stopsoffset, int64_t start, int64_t stop, int64_t step) awkward_listarray_getitem_next_range_carrylength @@ -1256,11 +1256,11 @@ - + - - void - void awkward_listarray32_getitem_next_range_carrylength + + Error + Error awkward_listarray32_getitem_next_range_carrylength (int64_t *carrylength, const int32_t *fromstarts, const int32_t *fromstops, int64_t lenstarts, int64_t startsoffset, int64_t stopsoffset, int64_t start, int64_t stop, int64_t step) awkward_listarray32_getitem_next_range_carrylength @@ -1305,11 +1305,11 @@ - + - - void - void awkward_listarray64_getitem_next_range_carrylength + + Error + Error awkward_listarray64_getitem_next_range_carrylength (int64_t *carrylength, const int64_t *fromstarts, const int64_t *fromstops, int64_t lenstarts, int64_t startsoffset, int64_t stopsoffset, int64_t start, int64_t stop, int64_t step) awkward_listarray64_getitem_next_range_carrylength @@ -1354,9 +1354,9 @@ - + - + typename C @@ -1365,8 +1365,8 @@ typename T - void - void awkward_listarray_getitem_next_range + Error + Error awkward_listarray_getitem_next_range (C *tooffsets, T *tocarry, const C *fromstarts, const C *fromstops, int64_t lenstarts, int64_t startsoffset, int64_t stopsoffset, int64_t start, int64_t stop, int64_t step) awkward_listarray_getitem_next_range @@ -1415,11 +1415,11 @@ - + - - void - void awkward_listarray32_getitem_next_range_64 + + Error + Error awkward_listarray32_getitem_next_range_64 (int32_t *tooffsets, int64_t *tocarry, const int32_t *fromstarts, const int32_t *fromstops, int64_t lenstarts, int64_t startsoffset, int64_t stopsoffset, int64_t start, int64_t stop, int64_t step) awkward_listarray32_getitem_next_range_64 @@ -1468,11 +1468,11 @@ - + - - void - void awkward_listarray64_getitem_next_range_64 + + Error + Error awkward_listarray64_getitem_next_range_64 (int64_t *tooffsets, int64_t *tocarry, const int64_t *fromstarts, const int64_t *fromstops, int64_t lenstarts, int64_t startsoffset, int64_t stopsoffset, int64_t start, int64_t stop, int64_t step) awkward_listarray64_getitem_next_range_64 @@ -1521,9 +1521,9 @@ - + - + typename C @@ -1532,8 +1532,8 @@ typename T - void - void awkward_listarray_getitem_next_range_counts + Error + Error awkward_listarray_getitem_next_range_counts (int64_t *total, const C *fromoffsets, int64_t lenstarts) awkward_listarray_getitem_next_range_counts @@ -1554,11 +1554,11 @@ - + - - void - void awkward_listarray32_getitem_next_range_counts_64 + + Error + Error awkward_listarray32_getitem_next_range_counts_64 (int64_t *total, const int32_t *fromoffsets, int64_t lenstarts) awkward_listarray32_getitem_next_range_counts_64 @@ -1579,11 +1579,11 @@ - + - - void - void awkward_listarray64_getitem_next_range_counts_64 + + Error + Error awkward_listarray64_getitem_next_range_counts_64 (int64_t *total, const int64_t *fromoffsets, int64_t lenstarts) awkward_listarray64_getitem_next_range_counts_64 @@ -1604,9 +1604,9 @@ - + - + typename C @@ -1615,8 +1615,8 @@ typename T - void - void awkward_listarray_getitem_next_range_spreadadvanced + Error + Error awkward_listarray_getitem_next_range_spreadadvanced (T *toadvanced, const T *fromadvanced, const C *fromoffsets, int64_t lenstarts) awkward_listarray_getitem_next_range_spreadadvanced @@ -1641,11 +1641,11 @@ - + - - void - void awkward_listarray32_getitem_next_range_spreadadvanced_64 + + Error + Error awkward_listarray32_getitem_next_range_spreadadvanced_64 (int64_t *toadvanced, const int64_t *fromadvanced, const int32_t *fromoffsets, int64_t lenstarts) awkward_listarray32_getitem_next_range_spreadadvanced_64 @@ -1670,11 +1670,11 @@ - + - - void - void awkward_listarray64_getitem_next_range_spreadadvanced_64 + + Error + Error awkward_listarray64_getitem_next_range_spreadadvanced_64 (int64_t *toadvanced, const int64_t *fromadvanced, const int64_t *fromoffsets, int64_t lenstarts) awkward_listarray64_getitem_next_range_spreadadvanced_64 @@ -1699,7 +1699,7 @@ - + @@ -1764,7 +1764,7 @@ - + Error @@ -1821,7 +1821,7 @@ - + Error @@ -1878,7 +1878,7 @@ - + @@ -1943,7 +1943,7 @@ - + Error @@ -2000,7 +2000,7 @@ - + Error @@ -2057,9 +2057,9 @@ - + - + typename C @@ -2068,9 +2068,9 @@ typename T - void - void awkward_listarray_getitem_carry - (C *tostarts, C *tostops, const C *fromstarts, const C *fromstops, const T *fromcarry, int64_t startsoffset, int64_t stopsoffset, int64_t lencarry) + Error + Error awkward_listarray_getitem_carry + (C *tostarts, C *tostops, const C *fromstarts, const C *fromstops, const T *fromcarry, int64_t startsoffset, int64_t stopsoffset, int64_t lenstarts, int64_t lencarry) awkward_listarray_getitem_carry C * @@ -2100,6 +2100,10 @@ int64_t stopsoffset + + int64_t + lenstarts + int64_t lencarry @@ -2110,12 +2114,12 @@ - + - - void - void awkward_listarray32_getitem_carry_64 - (int32_t *tostarts, int32_t *tostops, const int32_t *fromstarts, const int32_t *fromstops, const int64_t *fromcarry, int64_t startsoffset, int64_t stopsoffset, int64_t lencarry) + + Error + Error awkward_listarray32_getitem_carry_64 + (int32_t *tostarts, int32_t *tostops, const int32_t *fromstarts, const int32_t *fromstops, const int64_t *fromcarry, int64_t startsoffset, int64_t stopsoffset, int64_t lenstarts, int64_t lencarry) awkward_listarray32_getitem_carry_64 int32_t * @@ -2145,6 +2149,10 @@ int64_t stopsoffset + + int64_t + lenstarts + int64_t lencarry @@ -2155,12 +2163,12 @@ - + - - void - void awkward_listarray64_getitem_carry_64 - (int64_t *tostarts, int64_t *tostops, const int64_t *fromstarts, const int64_t *fromstops, const int64_t *fromcarry, int64_t startsoffset, int64_t stopsoffset, int64_t lencarry) + + Error + Error awkward_listarray64_getitem_carry_64 + (int64_t *tostarts, int64_t *tostops, const int64_t *fromstarts, const int64_t *fromstops, const int64_t *fromcarry, int64_t startsoffset, int64_t stopsoffset, int64_t lenstarts, int64_t lencarry) awkward_listarray64_getitem_carry_64 int64_t * @@ -2190,6 +2198,10 @@ int64_t stopsoffset + + int64_t + lenstarts + int64_t lencarry @@ -2200,7 +2212,7 @@ - + diff --git a/awkward1/signatures/identity_8cpp.xml b/awkward1/signatures/identity_8cpp.xml index 8d513c0160..a6cf3d5f96 100644 --- a/awkward1/signatures/identity_8cpp.xml +++ b/awkward1/signatures/identity_8cpp.xml @@ -15,14 +15,14 @@ - + typename T - void - void awkward_new_identity + Error + Error awkward_new_identity (T *toptr, int64_t length) awkward_new_identity @@ -39,11 +39,11 @@ - + - - void - void awkward_new_identity32 + + Error + Error awkward_new_identity32 (int32_t *toptr, int64_t length) awkward_new_identity32 @@ -60,11 +60,11 @@ - + - - void - void awkward_new_identity64 + + Error + Error awkward_new_identity64 (int64_t *toptr, int64_t length) awkward_new_identity64 @@ -81,11 +81,11 @@ - + - - void - void awkward_identity32_to_identity64 + + Error + Error awkward_identity32_to_identity64 (int64_t *toptr, const int32_t *fromptr, int64_t length) awkward_identity32_to_identity64 @@ -106,7 +106,7 @@ - + @@ -167,7 +167,7 @@ - + Error @@ -220,7 +220,7 @@ - + Error @@ -273,7 +273,7 @@ - + Error @@ -326,7 +326,7 @@ - + diff --git a/awkward1/signatures/libawkward_2util_8cpp.xml b/awkward1/signatures/libawkward_2util_8cpp.xml new file mode 100644 index 0000000000..08ce2c20bb --- /dev/null +++ b/awkward1/signatures/libawkward_2util_8cpp.xml @@ -0,0 +1,70 @@ + + + + util.cpp + cassert + sstream + awkward/util.h + awkward/Identity.h + + + + + + + + + + + + + + + + + + + + + + + + + + + awkward + awkward::util + + + void + void awkward::util::handle_error + (const Error &err, const std::string classname, const Identity *id) + handle_error + + const Error & + err + + + const std::string + classname + + + const Identity * + id + + + + + + + + + + + + + + + + + diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index e8cba85528..a95b6f42d5 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -90,19 +90,23 @@ def f3(q): assert awkward1.tolist(f3(listarray)) == [[8.8, 9.9], [4.4, 5.5], [0.0, 1.1], [6.6, 7.7]] -def test_current(): - content = awkward1.layout.NumpyArray(numpy.array([[0.0, 1.1], [2.2, 3.3], [4.4, 5.5], [6.6, 7.7], [8.8, 9.9]])) - listarray = awkward1.layout.ListArray64(awkward1.layout.Index64(numpy.array([0, 3, 3])), awkward1.layout.Index64(numpy.array([3, 3, 5])), content) - @numba.njit def f4(q): - return q[[2, 0, 0, -1], [1, -1, 0, 0], [0, 0, 0, 0]] + return q[[2, 0, 0, -1], [1, -1, 0, 0], -1] - # [[8.8, 9.9], [4.4, 5.5], [0.0, 1.1], [6.6, 7.7]] + assert awkward1.tolist(f4(listarray)) == [9.9, 5.5, 1.1, 7.7] - print(awkward1.tolist(f4(listarray))) + @numba.njit + def f5(q): + return q[[2, 0, 0, -1], [1, -1, 0, 0], -1:] + + assert awkward1.tolist(f5(listarray)) == [[9.9], [5.5], [1.1], [7.7]] + + @numba.njit + def f6(q): + return q[[2, 0, 0, -1], [1, -1, 0, 0], [0, 1, 0, 1]] - raise Exception + assert awkward1.tolist(f6(listarray)) == [8.8, 5.5, 0.0, 7.7] content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) From 51b70f5a11323ac240721a24358bf0f8faae15dc Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 15:29:44 -0500 Subject: [PATCH 19/36] Try to fix 32-bit error. --- awkward1/_numba/numpyarray.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/awkward1/_numba/numpyarray.py b/awkward1/_numba/numpyarray.py index 5fd86e4435..6aeebf6800 100644 --- a/awkward1/_numba/numpyarray.py +++ b/awkward1/_numba/numpyarray.py @@ -179,7 +179,7 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, lenself, skip = shapeunpacked[:2] lennext = builder.mul(lenself, skip) - carry = numba.targets.arrayobj.numpy_arange_1(context, builder, util.index64tpe(numba.int64), (lenself,)) + carry = numba.targets.arrayobj.numpy_arange_1(context, builder, util.index64tpe(numba.intp), (lenself,)) flathead = numba.targets.arrayobj.array_flatten(context, builder, util.index64tpe(headtpe), (headval,)) lenflathead = util.arraylen(context, builder, util.index64tpe, flathead) From 5dd1b0faa89f04bc05b17fd44c0b7837ea8bd0a4 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 17:13:08 -0500 Subject: [PATCH 20/36] ListArray.getitem_next(int) works. --- awkward1/_numba/listarray.py | 47 ++++++---- awkward1/_numba/listoffsetarray.py | 10 +-- .../test_PR015_finish_up_getitem_in_numba.py | 88 +++++++++++++++---- 3 files changed, 106 insertions(+), 39 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index 45e2053c7f..e1e728f66b 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -173,18 +173,11 @@ def lower_getitem_int(context, builder, sig, args): val, whereval = args proxyin = numba.cgutils.create_struct_proxy(tpe)(context, builder, value=val) - if tpe.bitwidth == 32: - starttpe, stoptpe = numba.int32, numba.int32 - elif tpe.bitwidth == 64: - starttpe, stoptpe = numba.int64, numba.int64 - else: - raise AssertionError("unrecognized bitwidth") - - start = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, starttpe(tpe.startstpe, wheretpe), (proxyin.starts, whereval)) - stop = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, stoptpe(tpe.stopstpe, wheretpe), (proxyin.stops, whereval)) + start = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, tpe.startstpe.dtype(tpe.startstpe, wheretpe), (proxyin.starts, whereval)) + stop = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, tpe.startstpe.dtype(tpe.stopstpe, wheretpe), (proxyin.stops, whereval)) proxyslice = numba.cgutils.create_struct_proxy(numba.types.slice2_type)(context, builder) - proxyslice.start = util.cast(context, builder, tpe, numba.intp, start) - proxyslice.stop = util.cast(context, builder, tpe, numba.intp, stop) + proxyslice.start = util.cast(context, builder, tpe.startstpe.dtype, numba.intp, start) + proxyslice.stop = util.cast(context, builder, tpe.stopstpe.dtype, numba.intp, stop) proxyslice.step = context.get_constant(numba.intp, 1) outtpe = tpe.contenttpe.getitem_range() @@ -242,6 +235,11 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, if len(wheretpe.types) == 0: return arrayval + headtpe = wheretpe.types[0] + tailtpe = numba.types.Tuple(wheretpe.types[1:]) + headval = numba.cgutils.unpack_tuple(builder, whereval)[0] + tailval = context.make_tuple(builder, tailtpe, numba.cgutils.unpack_tuple(builder, whereval)[1:]) + proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) lenstarts = util.arraylen(context, builder, arraytpe.startstpe, proxyin.starts, totpe=numba.int64) lenstops = util.arraylen(context, builder, arraytpe.stopstpe, proxyin.stops, totpe=numba.int64) @@ -250,13 +248,28 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, with builder.if_then(builder.icmp_signed("<", lenstops, lenstarts), likely=False): context.call_conv.return_user_exc(builder, ValueError, ("len(stops) < len(starts)",)) - headtpe = wheretpe.types[0] - tailtpe = numba.types.Tuple(wheretpe.types[1:]) - headval = numba.cgutils.unpack_tuple(builder, whereval)[0] - tailval = context.make_tuple(builder, tailtpe, numba.cgutils.unpack_tuple(builder, whereval)[1:]) - if isinstance(headtpe, numba.types.Integer): - raise NotImplementedError("ListArray.getitem_next(int)") + assert advanced is None + if arraytpe.bitwidth == 64: + kernel = cpu.kernels.awkward_listarray64_getitem_next_at_64 + elif arraytpe.bitwidth == 32: + kernel = cpu.kernels.awkward_listarray32_getitem_next_at_64 + else: + raise AssertionError("unrecognized bitwidth") + + nextcarry = util.newindex64(context, builder, numba.int64, lenstarts) + util.call(context, builder, kernel, + (util.arrayptr(context, builder, util.index64tpe, nextcarry), + util.arrayptr(context, builder, arraytpe.startstpe, proxyin.starts), + util.arrayptr(context, builder, arraytpe.stopstpe, proxyin.stops), + lenstarts, + context.get_constant(numba.int64, 0), + context.get_constant(numba.int64, 0), + util.cast(context, builder, headtpe, numba.int64, headval)), + "in {}, indexing error".format(arraytpe.shortname)) + nextcontenttpe = arraytpe.contenttpe.carry() + nextcontentval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) + return arraytpe.contenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, advanced) elif isinstance(headtpe, numba.types.SliceType): raise NotImplementedError("ListArray.getitem_next(slice)") diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index 378a8b8a07..f86680e343 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -247,6 +247,11 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, if len(wheretpe.types) == 0: return arrayval + headtpe = wheretpe.types[0] + tailtpe = numba.types.Tuple(wheretpe.types[1:]) + headval = numba.cgutils.unpack_tuple(builder, whereval)[0] + tailval = context.make_tuple(builder, tailtpe, numba.cgutils.unpack_tuple(builder, whereval)[1:]) + proxyin = numba.cgutils.create_struct_proxy(arraytpe)(context, builder, value=arrayval) lenoffsets = util.arraylen(context, builder, arraytpe.offsetstpe, proxyin.offsets, totpe=numba.int64) lenstarts = builder.sub(lenoffsets, context.get_constant(numba.int64, 1)) @@ -254,11 +259,6 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, starts, stops = starts_stops(context, builder, tpe.offsetstpe, proxyin.offsets, lenstarts, lenoffsets) - headtpe = wheretpe.types[0] - tailtpe = numba.types.Tuple(wheretpe.types[1:]) - headval = numba.cgutils.unpack_tuple(builder, whereval)[0] - tailval = context.make_tuple(builder, tailtpe, numba.cgutils.unpack_tuple(builder, whereval)[1:]) - if isinstance(headtpe, numba.types.Integer): raise NotImplementedError("ListArray.getitem_next(int)") diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index a95b6f42d5..f79375d379 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -108,20 +108,74 @@ def f6(q): assert awkward1.tolist(f6(listarray)) == [8.8, 5.5, 0.0, 7.7] -content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) -offsetsA = numpy.arange(0, 2*3*5 + 5, 5) -offsetsB = numpy.arange(0, 2*3 + 3, 3) -startsA, stopsA = offsetsA[:-1], offsetsA[1:] -startsB, stopsB = offsetsB[:-1], offsetsB[1:] - -listoffsetarrayA64 = awkward1.layout.ListOffsetArray64(awkward1.layout.Index64(offsetsA), content) -listoffsetarrayA32 = awkward1.layout.ListOffsetArray32(awkward1.layout.Index32(offsetsA), content) -listarrayA64 = awkward1.layout.ListArray64(awkward1.layout.Index64(startsA), awkward1.layout.Index64(stopsA), content) -listarrayA32 = awkward1.layout.ListArray32(awkward1.layout.Index32(startsA), awkward1.layout.Index32(stopsA), content) -modelA = numpy.arange(2*3*5*7).reshape(2*3, 5, 7) - -listoffsetarrayB64 = awkward1.layout.ListOffsetArray64(awkward1.layout.Index64(offsetsB), listoffsetarrayA64) -listoffsetarrayB32 = awkward1.layout.ListOffsetArray32(awkward1.layout.Index32(offsetsB), listoffsetarrayA64) -listarrayB64 = awkward1.layout.ListArray64(awkward1.layout.Index64(startsB), awkward1.layout.Index64(stopsB), listarrayA64) -listarrayB32 = awkward1.layout.ListArray32(awkward1.layout.Index32(startsB), awkward1.layout.Index32(stopsB), listarrayA64) -modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) +def test_deep_listarray32(): + content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) + offsetsA = numpy.arange(0, 2*3*5 + 5, 5) + offsetsB = numpy.arange(0, 2*3 + 3, 3) + startsA, stopsA = offsetsA[:-1], offsetsA[1:] + startsB, stopsB = offsetsB[:-1], offsetsB[1:] + + listarrayA32 = awkward1.layout.ListArray32(awkward1.layout.Index32(startsA), awkward1.layout.Index32(stopsA), content) + modelA = numpy.arange(2*3*5*7).reshape(2*3, 5, 7) + + listarrayB32 = awkward1.layout.ListArray32(awkward1.layout.Index32(startsB), awkward1.layout.Index32(stopsB), listarrayA32) + modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) + + @numba.njit + def f1(q): + return q[1, -1, 4] + + assert awkward1.tolist(f1(listarrayB32)) == awkward1.tolist(f1(modelB)) + + @numba.njit + def f2(q): + return q[1, -1, 4, -2] + + assert f2(listarrayB32) == f2(modelB) + +def test_deep_listoffsetarray32(): + content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) + offsetsA = numpy.arange(0, 2*3*5 + 5, 5) + offsetsB = numpy.arange(0, 2*3 + 3, 3) + + listoffsetarrayA32 = awkward1.layout.ListOffsetArray32(awkward1.layout.Index32(offsetsA), content) + modelA = numpy.arange(2*3*5*7).reshape(2*3, 5, 7) + + listoffsetarrayB32 = awkward1.layout.ListOffsetArray32(awkward1.layout.Index32(offsetsB), listoffsetarrayA32) + modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) + +def test_deep_listarray64(): + content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) + offsetsA = numpy.arange(0, 2*3*5 + 5, 5) + offsetsB = numpy.arange(0, 2*3 + 3, 3) + startsA, stopsA = offsetsA[:-1], offsetsA[1:] + startsB, stopsB = offsetsB[:-1], offsetsB[1:] + + listarrayA64 = awkward1.layout.ListArray64(awkward1.layout.Index64(startsA), awkward1.layout.Index64(stopsA), content) + modelA = numpy.arange(2*3*5*7).reshape(2*3, 5, 7) + + listarrayB64 = awkward1.layout.ListArray64(awkward1.layout.Index64(startsB), awkward1.layout.Index64(stopsB), listarrayA64) + modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) + + @numba.njit + def f1(q): + return q[1, -1, 4] + + assert awkward1.tolist(f1(listarrayB64)) == awkward1.tolist(f1(modelB)) + + @numba.njit + def f2(q): + return q[1, -1, 4, -2] + + assert f2(listarrayB64) == f2(modelB) + +def test_deep_listoffsetarray64(): + content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) + offsetsA = numpy.arange(0, 2*3*5 + 5, 5) + offsetsB = numpy.arange(0, 2*3 + 3, 3) + + listoffsetarrayA64 = awkward1.layout.ListOffsetArray64(awkward1.layout.Index64(offsetsA), content) + modelA = numpy.arange(2*3*5*7).reshape(2*3, 5, 7) + + listoffsetarrayB64 = awkward1.layout.ListOffsetArray64(awkward1.layout.Index64(offsetsB), listoffsetarrayA64) + modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) From f522bc2de0f7996ee66020b524e42037d4337326 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 17:33:42 -0500 Subject: [PATCH 21/36] [skip ci] in progress. --- awkward1/_numba/listarray.py | 2 +- awkward1/_numba/listoffsetarray.py | 44 ++++++++++++++----- .../test_PR015_finish_up_getitem_in_numba.py | 22 ++++++++++ 3 files changed, 57 insertions(+), 11 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index e1e728f66b..66ecad18fe 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -269,7 +269,7 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, "in {}, indexing error".format(arraytpe.shortname)) nextcontenttpe = arraytpe.contenttpe.carry() nextcontentval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) - return arraytpe.contenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, advanced) + return nextcontenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, advanced) elif isinstance(headtpe, numba.types.SliceType): raise NotImplementedError("ListArray.getitem_next(slice)") diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index f86680e343..bccf79ebaa 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -36,8 +36,8 @@ def getitem_range(self): return self def getitem_tuple(self, wheretpe): - import awkward._numba.listarray - nexttpe = awkward._numba.listarray.ListArrayType(util.index64tpe, util.index64tpe, self, numba.none) + import awkward1._numba.listarray + nexttpe = awkward1._numba.listarray.ListArrayType(util.index64tpe, util.index64tpe, self, numba.none) out = nexttpe.getitem_next(wheretpe, False) return out.getitem_int() @@ -66,7 +66,7 @@ def getitem_next(self, wheretpe, isadvanced): def carry(self): import awkward1._numba.listarray - return awkward1._numba.listarray.ListArrayType(self.offsettpe, self.offsettpe, self.contenttpe, self.idtpe) + return awkward1._numba.listarray.ListArrayType(self.offsetstpe, self.offsetstpe, self.contenttpe, self.idtpe) @property def lower_len(self): @@ -212,7 +212,11 @@ def lower_getitem_tuple(context, builder, sig, args): outtpe = nexttpe.getitem_next(wheretpe, False) outval = nexttpe.lower_getitem_next(context, builder, nexttpe, wheretpe, nextval, whereval, None) - return outtpe.lower_getitem_int(context, builder, rettpe(outtpe, numba.int64), (outval, conext.get_constant(numba.int64, 0))) + print("outtpe", outtpe) + print("rettpe", rettpe) + print("outtpe.contenttpe", outtpe.contenttpe) + + return outtpe.lower_getitem_int(context, builder, rettpe(outtpe, numba.int64), (outval, context.get_constant(numba.int64, 0))) @numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.Array) @numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.List) @@ -260,16 +264,36 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, starts, stops = starts_stops(context, builder, tpe.offsetstpe, proxyin.offsets, lenstarts, lenoffsets) if isinstance(headtpe, numba.types.Integer): - raise NotImplementedError("ListArray.getitem_next(int)") + assert advanced is None + if arraytpe.bitwidth == 64: + kernel = cpu.kernels.awkward_listarray64_getitem_next_at_64 + elif arraytpe.bitwidth == 32: + kernel = cpu.kernels.awkward_listarray32_getitem_next_at_64 + else: + raise AssertionError("unrecognized bitwidth") + + nextcarry = util.newindex64(context, builder, numba.int64, lenstarts) + util.call(context, builder, kernel, + (util.arrayptr(context, builder, util.index64tpe, nextcarry), + util.arrayptr(context, builder, arraytpe.startstpe, starts), + util.arrayptr(context, builder, arraytpe.stopstpe, stops), + lenstarts, + context.get_constant(numba.int64, 0), + context.get_constant(numba.int64, 0), + util.cast(context, builder, headtpe, numba.int64, headval)), + "in {}, indexing error".format(arraytpe.shortname)) + nextcontenttpe = arraytpe.contenttpe.carry() + nextcontentval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) + return nextcontenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, advanced) elif isinstance(headtpe, numba.types.SliceType): - raise NotImplementedError("ListArray.getitem_next(slice)") + raise NotImplementedError("ListOffsetArray.getitem_next(slice)") elif isinstance(headtpe, numba.types.EllipsisType): - raise NotImplementedError("ListArray.getitem_next(ellipsis)") + raise NotImplementedError("ListOffsetArray.getitem_next(ellipsis)") elif isinstance(headtpe, type(numba.typeof(numpy.newaxis))): - raise NotImplementedError("ListArray.getitem_next(newaxis)") + raise NotImplementedError("ListOffsetArray.getitem_next(newaxis)") elif isinstance(headtpe, numba.types.Array) and advanced is None: if headtpe.ndim != 1: @@ -324,9 +348,9 @@ def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): lenoffsets = util.arraylen(context, builder, arraytpe.offsetstpe, proxyin.offsets, totpe=numba.int64) lenstarts = builder.sub(lenoffsets, context.get_constant(numba.int64, 1)) - starts, stops = starts_stops(context, builder, tpe.offsetstpe, proxyin.offsets, lenstarts, lenoffsets) + starts, stops = starts_stops(context, builder, arraytpe.offsetstpe, proxyin.offsets, lenstarts, lenoffsets) - proxyout = numba.cgutils.create_struct_proxy(awkward1._numba.listarray.ListArray(arraytpe.offsetstpe, arraytpe.offsetstpe, arraytpe.contenttpe, arraytpe.idtpe))(context, builder) + proxyout = numba.cgutils.create_struct_proxy(awkward1._numba.listarray.ListArrayType(arraytpe.offsetstpe, arraytpe.offsetstpe, arraytpe.contenttpe, arraytpe.idtpe))(context, builder) proxyout.starts = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (starts, carryval)) proxyout.stops = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (stops, carryval)) proxyout.content = proxyin.content diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index f79375d379..0e9a3dbe9a 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -179,3 +179,25 @@ def test_deep_listoffsetarray64(): listoffsetarrayB64 = awkward1.layout.ListOffsetArray64(awkward1.layout.Index64(offsetsB), listoffsetarrayA64) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) + + @numba.njit + def f1(q): + return q[1, 1] + + tmp = f1(listoffsetarrayB64) + print(tmp) + print(awkward1.tolist(tmp)) + print(awkward1.tolist(f1(modelB))) + raise Exception + + # @numba.njit + # def f1(q): + # return q[1, -1, 4] + # + # assert awkward1.tolist(f1(listoffsetarrayB64)) == awkward1.tolist(f1(modelB)) + + # @numba.njit + # def f2(q): + # return q[1, -1, 4, -2] + # + # assert f2(listoffsetarrayB64) == f2(modelB) From 332a5c0b1bae7ff6a34239ad946567533eeb7220 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 17:52:13 -0500 Subject: [PATCH 22/36] [skip ci] fixed a carry-issue in type, but there's badly formed LLVM IR. --- awkward1/_numba/listarray.py | 4 +++- awkward1/_numba/listoffsetarray.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index 66ecad18fe..a3977e03da 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -53,7 +53,7 @@ def getitem_next(self, wheretpe, isadvanced): headtpe = wheretpe.types[0] tailtpe = numba.types.Tuple(wheretpe.types[1:]) if isinstance(headtpe, numba.types.Integer): - return self.contenttpe.getitem_next(tailtpe, isadvanced) + return self.contenttpe.carry().getitem_next(tailtpe, isadvanced) elif isinstance(headtpe, numba.types.SliceType): return awkward1._numba.listoffsetarray.ListOffsetArrayType(self.startstpe, self.contenttpe.getitem_next(tailtpe, isadvanced), self.idtpe) elif isinstance(headtpe, numba.types.EllipsisType): @@ -249,6 +249,8 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, context.call_conv.return_user_exc(builder, ValueError, ("len(stops) < len(starts)",)) if isinstance(headtpe, numba.types.Integer): + print("ListArray.getitem_next(int)") + assert advanced is None if arraytpe.bitwidth == 64: kernel = cpu.kernels.awkward_listarray64_getitem_next_at_64 diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index bccf79ebaa..a65a886e44 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -47,7 +47,7 @@ def getitem_next(self, wheretpe, isadvanced): headtpe = wheretpe.types[0] tailtpe = numba.types.Tuple(wheretpe.types[1:]) if isinstance(headtpe, numba.types.Integer): - return self.contenttpe.getitem_next(tailtpe, isadvanced) + return self.contenttpe.carry().getitem_next(tailtpe, isadvanced) elif isinstance(headtpe, numba.types.SliceType): return ListOffsetArrayType(self.offsetstpe, self.contenttpe.getitem_next(tailtpe, isadvanced), self.idtpe) elif isinstance(headtpe, numba.types.EllipsisType): @@ -264,6 +264,8 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, starts, stops = starts_stops(context, builder, tpe.offsetstpe, proxyin.offsets, lenstarts, lenoffsets) if isinstance(headtpe, numba.types.Integer): + print("ListOffsetArray.getitem_next(int)") + assert advanced is None if arraytpe.bitwidth == 64: kernel = cpu.kernels.awkward_listarray64_getitem_next_at_64 From be5b5016b987d0965fdc4de58e666554b524725b Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 19:03:39 -0500 Subject: [PATCH 23/36] ListArray32/64 and ListOffsetArray32/64 are all at the same level: getitem_next(int*) works. --- awkward1/_numba/listarray.py | 2 - awkward1/_numba/listoffsetarray.py | 18 ++++----- .../test_PR015_finish_up_getitem_in_numba.py | 40 ++++++++++--------- 3 files changed, 28 insertions(+), 32 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index a3977e03da..c1136c99d2 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -249,8 +249,6 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, context.call_conv.return_user_exc(builder, ValueError, ("len(stops) < len(starts)",)) if isinstance(headtpe, numba.types.Integer): - print("ListArray.getitem_next(int)") - assert advanced is None if arraytpe.bitwidth == 64: kernel = cpu.kernels.awkward_listarray64_getitem_next_at_64 diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index a65a886e44..57c614fd40 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -211,11 +211,6 @@ def lower_getitem_tuple(context, builder, sig, args): outtpe = nexttpe.getitem_next(wheretpe, False) outval = nexttpe.lower_getitem_next(context, builder, nexttpe, wheretpe, nextval, whereval, None) - - print("outtpe", outtpe) - print("rettpe", rettpe) - print("outtpe.contenttpe", outtpe.contenttpe) - return outtpe.lower_getitem_int(context, builder, rettpe(outtpe, numba.int64), (outval, context.get_constant(numba.int64, 0))) @numba.extending.lower_builtin(operator.getitem, ListOffsetArrayType, numba.types.Array) @@ -233,16 +228,18 @@ def lower_getitem_other(context, builder, sig, args): def starts_stops(context, builder, offsetstpe, offsetsval, lenstarts, lenoffsets): proxyslicestarts = numba.cgutils.create_struct_proxy(numba.types.slice2_type)(context, builder) proxyslicestarts.start = context.get_constant(numba.intp, 0) - proxyslicestarts.stop = context.get_constant(numba.intp, lenstarts) + proxyslicestarts.stop = lenstarts proxyslicestarts.step = context.get_constant(numba.intp, 1) - starts = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, offsetstpe(offsetstpe, numba.types.slice2_type), (offsetsval, proxyslicestarts._getvalue())) + slicestarts = proxyslicestarts._getvalue() proxyslicestops = numba.cgutils.create_struct_proxy(numba.types.slice2_type)(context, builder) proxyslicestops.start = context.get_constant(numba.intp, 1) - proxyslicestops.stop = context.get_constant(numba.intp, lenoffsets) + proxyslicestops.stop = lenoffsets proxyslicestops.step = context.get_constant(numba.intp, 1) - stops = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, offsetstpe(offsetstpe, numba.types.slice2_type), (offsetsval, proxyslicestops._getvalue())) + slicestops = proxyslicestops._getvalue() + starts = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, offsetstpe(offsetstpe, numba.types.slice2_type), (offsetsval, slicestarts)) + stops = numba.targets.arrayobj.getitem_arraynd_intp(context, builder, offsetstpe(offsetstpe, numba.types.slice2_type), (offsetsval, slicestops)) return starts, stops def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, advanced): @@ -264,8 +261,6 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, starts, stops = starts_stops(context, builder, tpe.offsetstpe, proxyin.offsets, lenstarts, lenoffsets) if isinstance(headtpe, numba.types.Integer): - print("ListOffsetArray.getitem_next(int)") - assert advanced is None if arraytpe.bitwidth == 64: kernel = cpu.kernels.awkward_listarray64_getitem_next_at_64 @@ -355,6 +350,7 @@ def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): proxyout = numba.cgutils.create_struct_proxy(awkward1._numba.listarray.ListArrayType(arraytpe.offsetstpe, arraytpe.offsetstpe, arraytpe.contenttpe, arraytpe.idtpe))(context, builder) proxyout.starts = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (starts, carryval)) proxyout.stops = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (stops, carryval)) + proxyout.content = proxyin.content if arraytpe.idtpe != numba.none: proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, arraytpe.idtpe, carrytpe, proxyin.id, carryval) diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index 0e9a3dbe9a..9415c9b26a 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -144,6 +144,18 @@ def test_deep_listoffsetarray32(): listoffsetarrayB32 = awkward1.layout.ListOffsetArray32(awkward1.layout.Index32(offsetsB), listoffsetarrayA32) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) + @numba.njit + def f1(q): + return q[1, -1, 4] + + assert awkward1.tolist(f1(listoffsetarrayB32)) == awkward1.tolist(f1(modelB)) + + @numba.njit + def f2(q): + return q[1, -1, 4, -2] + + assert f2(listoffsetarrayB32) == f2(modelB) + def test_deep_listarray64(): content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) @@ -182,22 +194,12 @@ def test_deep_listoffsetarray64(): @numba.njit def f1(q): - return q[1, 1] - - tmp = f1(listoffsetarrayB64) - print(tmp) - print(awkward1.tolist(tmp)) - print(awkward1.tolist(f1(modelB))) - raise Exception - - # @numba.njit - # def f1(q): - # return q[1, -1, 4] - # - # assert awkward1.tolist(f1(listoffsetarrayB64)) == awkward1.tolist(f1(modelB)) - - # @numba.njit - # def f2(q): - # return q[1, -1, 4, -2] - # - # assert f2(listoffsetarrayB64) == f2(modelB) + return q[1, -1, 4] + + assert awkward1.tolist(f1(listoffsetarrayB64)) == awkward1.tolist(f1(modelB)) + + @numba.njit + def f2(q): + return q[1, -1, 4, -2] + + assert f2(listoffsetarrayB64) == f2(modelB) From 429b9513a43806e518d84c45d2f22da546a0746e Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 19:54:24 -0500 Subject: [PATCH 24/36] Probably fixed 32-bit issue. --- awkward1/_numba/listoffsetarray.py | 4 ++-- awkward1/_numba/util.py | 11 +++++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index 57c614fd40..1779d2ebed 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -228,13 +228,13 @@ def lower_getitem_other(context, builder, sig, args): def starts_stops(context, builder, offsetstpe, offsetsval, lenstarts, lenoffsets): proxyslicestarts = numba.cgutils.create_struct_proxy(numba.types.slice2_type)(context, builder) proxyslicestarts.start = context.get_constant(numba.intp, 0) - proxyslicestarts.stop = lenstarts + proxyslicestarts.stop = util.cast(context, builder, lenstarts.type, numba.intp, lenstarts) proxyslicestarts.step = context.get_constant(numba.intp, 1) slicestarts = proxyslicestarts._getvalue() proxyslicestops = numba.cgutils.create_struct_proxy(numba.types.slice2_type)(context, builder) proxyslicestops.start = context.get_constant(numba.intp, 1) - proxyslicestops.stop = lenoffsets + proxyslicestops.stop = util.cast(context, builder, lenoffsets.type, numba.intp, lenoffsets) proxyslicestops.step = context.get_constant(numba.intp, 1) slicestops = proxyslicestops._getvalue() diff --git a/awkward1/_numba/util.py b/awkward1/_numba/util.py index 6004f1b23b..8ef1b1fa7e 100644 --- a/awkward1/_numba/util.py +++ b/awkward1/_numba/util.py @@ -23,6 +23,17 @@ def debug(context, builder, *args): """, globals()) def cast(context, builder, fromtpe, totpe, val): + if isinstance(fromtpe, llvmlite.ir.types.IntType): + if fromtpe.width == 8: + fromtpe = numba.int8 + elif fromtpe.width == 16: + fromtpe = numba.int16 + elif fromtpe.width == 32: + fromtpe = numba.int32 + elif fromtpe.width == 64: + fromtpe = numba.int64 + else: + raise AssertionError("unrecognized bitwidth") if fromtpe.bitwidth < totpe.bitwidth: return builder.zext(val, context.get_value_type(totpe)) elif fromtpe.bitwidth > totpe.bitwidth: From fa03bf67e6b23e8cb45edd3e6e0a6e4521924e78 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 20:12:35 -0500 Subject: [PATCH 25/36] Try the test without negative-correction on 32-bit. --- tests/test_PR015_finish_up_getitem_in_numba.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index 9415c9b26a..f7d3ec07d9 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -194,7 +194,7 @@ def test_deep_listoffsetarray64(): @numba.njit def f1(q): - return q[1, -1, 4] + return q[1, 2, 4] assert awkward1.tolist(f1(listoffsetarrayB64)) == awkward1.tolist(f1(modelB)) From 10e6f0f52daca390a9db02728c3c57a31d82b41d Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 20:24:59 -0500 Subject: [PATCH 26/36] More diagnostics for 32-bit. --- awkward1/_numba/listarray.py | 4 ++++ awkward1/_numba/listoffsetarray.py | 8 ++++++++ tests/test_PR015_finish_up_getitem_in_numba.py | 6 +++--- 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index c1136c99d2..608aab508c 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -257,6 +257,10 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, else: raise AssertionError("unrecognized bitwidth") + util.debug(context, builder, headtpe, headval) + util.debug(context, builder, arraytpe.startstpe, proxyin.starts) + util.debug(context, builder, arraytpe.stopstpe, proxyin.stops) + nextcarry = util.newindex64(context, builder, numba.int64, lenstarts) util.call(context, builder, kernel, (util.arrayptr(context, builder, util.index64tpe, nextcarry), diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index 1779d2ebed..e5887b7701 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -269,6 +269,9 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, else: raise AssertionError("unrecognized bitwidth") + util.debug(context, builder, headtpe, headval) + util.debug(context, builder, arraytpe.offsetstpe, proxyin.offsets) + nextcarry = util.newindex64(context, builder, numba.int64, lenstarts) util.call(context, builder, kernel, (util.arrayptr(context, builder, util.index64tpe, nextcarry), @@ -347,6 +350,11 @@ def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): starts, stops = starts_stops(context, builder, arraytpe.offsetstpe, proxyin.offsets, lenstarts, lenoffsets) + util.debug(context, builder, numba.int64, lenoffsets, numba.int64, lenstarts) + util.debug(context, builder, arraytpe.offsetstpe, proxyin.offsets) + util.debug(context, builder, arraytpe.offsetstpe, starts) + util.debug(context, builder, arraytpe.offsetstpe, stops) + proxyout = numba.cgutils.create_struct_proxy(awkward1._numba.listarray.ListArrayType(arraytpe.offsetstpe, arraytpe.offsetstpe, arraytpe.contenttpe, arraytpe.idtpe))(context, builder) proxyout.starts = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (starts, carryval)) proxyout.stops = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (stops, carryval)) diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index f7d3ec07d9..3669645750 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -123,7 +123,7 @@ def test_deep_listarray32(): @numba.njit def f1(q): - return q[1, -1, 4] + return q[1, 2, 4] assert awkward1.tolist(f1(listarrayB32)) == awkward1.tolist(f1(modelB)) @@ -146,7 +146,7 @@ def test_deep_listoffsetarray32(): @numba.njit def f1(q): - return q[1, -1, 4] + return q[1, 2, 4] assert awkward1.tolist(f1(listoffsetarrayB32)) == awkward1.tolist(f1(modelB)) @@ -171,7 +171,7 @@ def test_deep_listarray64(): @numba.njit def f1(q): - return q[1, -1, 4] + return q[1, 2, 4] assert awkward1.tolist(f1(listarrayB64)) == awkward1.tolist(f1(modelB)) From f48e4b475c262e28c84534e4e2e5e7ba908c843c Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 20:42:10 -0500 Subject: [PATCH 27/36] Yet more diagnostics for 32-bit. --- src/cpu-kernels/getitem.cpp | 23 +++++++++++++++++++ .../test_PR015_finish_up_getitem_in_numba.py | 1 + 2 files changed, 24 insertions(+) diff --git a/src/cpu-kernels/getitem.cpp b/src/cpu-kernels/getitem.cpp index b9ec3e8e2a..6f026990e3 100644 --- a/src/cpu-kernels/getitem.cpp +++ b/src/cpu-kernels/getitem.cpp @@ -211,17 +211,40 @@ Error awkward_numpyarray_getitem_next_array_advanced_64(int64_t* nextcarryptr, c template Error awkward_listarray_getitem_next_at(T* tocarry, const C* fromstarts, const C* fromstops, int64_t lenstarts, int64_t startsoffset, int64_t stopsoffset, int64_t at) { + std::cout << "ENTER " << lenstarts << " " << startsoffset << " " << stopsoffset << " at " << at << std::endl; + for (int64_t i = 0; i < lenstarts; i++) { + std::cout << "i " << i << " start " << fromstarts[startsoffset + i] << " stop " << fromstops[stopsoffset + i] << std::endl; + int64_t length = fromstops[stopsoffset + i] - fromstarts[startsoffset + i]; + + std::cout << "length " << length << std::endl; + int64_t regular_at = at; + + std::cout << "regular_at " << regular_at << std::endl; + if (regular_at < 0) { regular_at += length; } + + std::cout << "regular_at " << regular_at << std::endl; + + std::cout << "0 <= regular_at " << (0 <= regular_at) << std::endl; + std::cout << "regular_at < length " << (regular_at < length) << std::endl; + std::cout << "(0 <= regular_at && regular_at < length) " << ((0 <= regular_at && regular_at < length)) << std::endl; + std::cout << "!(0 <= regular_at && regular_at < length) " << (!(0 <= regular_at && regular_at < length)) << std::endl; + if (!(0 <= regular_at && regular_at < length)) { + std::cout << "FAILURE" << std::endl; + return failure("index out of range", i, at); } tocarry[i] = fromstarts[startsoffset + i] + regular_at; + + std::cout << "tocarry[i] " << (tocarry[i]) << std::endl; } + std::cout << "SUCCESS" << std::endl; return success(); } Error awkward_listarray32_getitem_next_at_64(int64_t* tocarry, const int32_t* fromstarts, const int32_t* fromstops, int64_t lenstarts, int64_t startsoffset, int64_t stopsoffset, int64_t at) { diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index 3669645750..b997ee7b9b 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -180,6 +180,7 @@ def f2(q): return q[1, -1, 4, -2] assert f2(listarrayB64) == f2(modelB) + raise Exception def test_deep_listoffsetarray64(): content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) From 61bd975261cca2e8d9cd93629e5ae766a7fab9b5 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 21:02:32 -0500 Subject: [PATCH 28/36] Yet yet more diagnostics for 32-bit. --- awkward1/_numba/listarray.py | 14 +++++++++-- awkward1/_numba/util.py | 12 ++++++++++ src/cpu-kernels/getitem.cpp | 23 ------------------- .../test_PR015_finish_up_getitem_in_numba.py | 17 +++++++++++++- 4 files changed, 40 insertions(+), 26 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index 608aab508c..4ae06ff027 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -51,6 +51,9 @@ def getitem_next(self, wheretpe, isadvanced): if len(wheretpe.types) == 0: return self headtpe = wheretpe.types[0] + + print("getitem_next", headtpe) + tailtpe = numba.types.Tuple(wheretpe.types[1:]) if isinstance(headtpe, numba.types.Integer): return self.contenttpe.carry().getitem_next(tailtpe, isadvanced) @@ -209,9 +212,13 @@ def lower_getitem_tuple(context, builder, sig, args): rettpe, (arraytpe, wheretpe) = sig.return_type, sig.args arrayval, whereval = args + print("lower_getitem_tuple BEFORE", wheretpe) + wheretpe, whereval = util.preprocess_slicetuple(context, builder, wheretpe, whereval) nexttpe, nextval = util.wrap_for_slicetuple(context, builder, arraytpe, arrayval) + print("lower_getitem_tuple AFTER", wheretpe) + outtpe = nexttpe.getitem_next(wheretpe, False) outval = nexttpe.lower_getitem_next(context, builder, nexttpe, wheretpe, nextval, whereval, None) @@ -257,10 +264,13 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, else: raise AssertionError("unrecognized bitwidth") - util.debug(context, builder, headtpe, headval) + tmp = util.cast(context, builder, headtpe, numba.int64, headval) + util.debug(context, builder, headtpe, headval, numba.int64, tmp) util.debug(context, builder, arraytpe.startstpe, proxyin.starts) util.debug(context, builder, arraytpe.stopstpe, proxyin.stops) + print("headtpe", headtpe) + nextcarry = util.newindex64(context, builder, numba.int64, lenstarts) util.call(context, builder, kernel, (util.arrayptr(context, builder, util.index64tpe, nextcarry), @@ -269,7 +279,7 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, lenstarts, context.get_constant(numba.int64, 0), context.get_constant(numba.int64, 0), - util.cast(context, builder, headtpe, numba.int64, headval)), + tmp), "in {}, indexing error".format(arraytpe.shortname)) nextcontenttpe = arraytpe.contenttpe.carry() nextcontentval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) diff --git a/awkward1/_numba/util.py b/awkward1/_numba/util.py index 8ef1b1fa7e..4e0ec5edf1 100644 --- a/awkward1/_numba/util.py +++ b/awkward1/_numba/util.py @@ -141,12 +141,15 @@ def typing_broadcast_arrays(arrays): @numba.generated_jit(nopython=True) def regularize_slice(arrays): if not isinstance(arrays, numba.types.BaseTuple) and isinstance(arrays, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(arrays.dtype, numba.types.Boolean): + print("regularize_slice A", arrays) return lambda arrays: numpy.nonzero(arrays) elif not isinstance(arrays, numba.types.BaseTuple) or not any(isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) for t in arrays.types): + print("regularize_slice B", arrays) return lambda arrays: arrays else: + print("regularize_slice C", arrays) code = "def impl(arrays):\n" indexes = [] for i, t in enumerate(arrays.types): @@ -167,12 +170,15 @@ def regularize_slice(arrays): def typing_regularize_slice(arrays): out = () if not isinstance(arrays, numba.types.BaseTuple) and isinstance(arrays, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(arrays.dtype, numba.types.Boolean): + print("typing_regularize_slice A", arrays) return numba.types.Tuple(arrays.ndims*(numba.types.Array(numba.int64, 1, "C"),)) elif not isinstance(arrays, numba.types.BaseTuple) or not any(isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) for t in arrays.types): + print("typing_regularize_slice B", arrays) return arrays else: + print("typing_regularize_slice C", arrays) for t in arrays.types: if isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(t.dtype, numba.types.Boolean): out = out + t.ndims*(numba.types.Array(numba.int64, 1, "C"),) @@ -185,16 +191,22 @@ def typing_regularize_slice(arrays): return numba.types.Tuple(out) def preprocess_slicetuple(context, builder, wheretpe, whereval): + print("wheretpe1", wheretpe) + wheretpe2 = typing_regularize_slice(wheretpe) regularize_slice.compile(wheretpe2(wheretpe)) cres = regularize_slice.overloads[(wheretpe,)] whereval2 = context.call_internal(builder, cres.fndesc, wheretpe2(wheretpe), (whereval,)) + print("wheretpe2", wheretpe2) + wheretpe3 = typing_broadcast_arrays(wheretpe2) broadcast_arrays.compile(wheretpe3(wheretpe2)) cres2 = broadcast_arrays.overloads[(wheretpe2,)] whereval3 = context.call_internal(builder, cres2.fndesc, wheretpe3(wheretpe2), (whereval2,)) + print("wheretpe3", wheretpe3) + return wheretpe3, whereval3 def wrap_for_slicetuple(context, builder, arraytpe, arrayval): diff --git a/src/cpu-kernels/getitem.cpp b/src/cpu-kernels/getitem.cpp index 6f026990e3..b9ec3e8e2a 100644 --- a/src/cpu-kernels/getitem.cpp +++ b/src/cpu-kernels/getitem.cpp @@ -211,40 +211,17 @@ Error awkward_numpyarray_getitem_next_array_advanced_64(int64_t* nextcarryptr, c template Error awkward_listarray_getitem_next_at(T* tocarry, const C* fromstarts, const C* fromstops, int64_t lenstarts, int64_t startsoffset, int64_t stopsoffset, int64_t at) { - std::cout << "ENTER " << lenstarts << " " << startsoffset << " " << stopsoffset << " at " << at << std::endl; - for (int64_t i = 0; i < lenstarts; i++) { - std::cout << "i " << i << " start " << fromstarts[startsoffset + i] << " stop " << fromstops[stopsoffset + i] << std::endl; - int64_t length = fromstops[stopsoffset + i] - fromstarts[startsoffset + i]; - - std::cout << "length " << length << std::endl; - int64_t regular_at = at; - - std::cout << "regular_at " << regular_at << std::endl; - if (regular_at < 0) { regular_at += length; } - - std::cout << "regular_at " << regular_at << std::endl; - - std::cout << "0 <= regular_at " << (0 <= regular_at) << std::endl; - std::cout << "regular_at < length " << (regular_at < length) << std::endl; - std::cout << "(0 <= regular_at && regular_at < length) " << ((0 <= regular_at && regular_at < length)) << std::endl; - std::cout << "!(0 <= regular_at && regular_at < length) " << (!(0 <= regular_at && regular_at < length)) << std::endl; - if (!(0 <= regular_at && regular_at < length)) { - std::cout << "FAILURE" << std::endl; - return failure("index out of range", i, at); } tocarry[i] = fromstarts[startsoffset + i] + regular_at; - - std::cout << "tocarry[i] " << (tocarry[i]) << std::endl; } - std::cout << "SUCCESS" << std::endl; return success(); } Error awkward_listarray32_getitem_next_at_64(int64_t* tocarry, const int32_t* fromstarts, const int32_t* fromstops, int64_t lenstarts, int64_t startsoffset, int64_t stopsoffset, int64_t at) { diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index b997ee7b9b..88a1592cb6 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -121,12 +121,16 @@ def test_deep_listarray32(): listarrayB32 = awkward1.layout.ListArray32(awkward1.layout.Index32(startsB), awkward1.layout.Index32(stopsB), listarrayA32) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) + print("============================================") + @numba.njit def f1(q): return q[1, 2, 4] assert awkward1.tolist(f1(listarrayB32)) == awkward1.tolist(f1(modelB)) + print("--------------------------------------------") + @numba.njit def f2(q): return q[1, -1, 4, -2] @@ -144,12 +148,16 @@ def test_deep_listoffsetarray32(): listoffsetarrayB32 = awkward1.layout.ListOffsetArray32(awkward1.layout.Index32(offsetsB), listoffsetarrayA32) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) + print("============================================") + @numba.njit def f1(q): return q[1, 2, 4] assert awkward1.tolist(f1(listoffsetarrayB32)) == awkward1.tolist(f1(modelB)) + print("--------------------------------------------") + @numba.njit def f2(q): return q[1, -1, 4, -2] @@ -169,18 +177,21 @@ def test_deep_listarray64(): listarrayB64 = awkward1.layout.ListArray64(awkward1.layout.Index64(startsB), awkward1.layout.Index64(stopsB), listarrayA64) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) + print("============================================") + @numba.njit def f1(q): return q[1, 2, 4] assert awkward1.tolist(f1(listarrayB64)) == awkward1.tolist(f1(modelB)) + print("--------------------------------------------") + @numba.njit def f2(q): return q[1, -1, 4, -2] assert f2(listarrayB64) == f2(modelB) - raise Exception def test_deep_listoffsetarray64(): content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) @@ -193,12 +204,16 @@ def test_deep_listoffsetarray64(): listoffsetarrayB64 = awkward1.layout.ListOffsetArray64(awkward1.layout.Index64(offsetsB), listoffsetarrayA64) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) + print("============================================") + @numba.njit def f1(q): return q[1, 2, 4] assert awkward1.tolist(f1(listoffsetarrayB64)) == awkward1.tolist(f1(modelB)) + print("--------------------------------------------") + @numba.njit def f2(q): return q[1, -1, 4, -2] From bb2dae0a27f08ec6b038f17d666f18609ec95f96 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 21:11:57 -0500 Subject: [PATCH 29/36] Replace 'zext' with 'sext'. --- awkward1/_numba/listarray.py | 9 --------- awkward1/_numba/util.py | 2 +- 2 files changed, 1 insertion(+), 10 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index 4ae06ff027..5934297f45 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -51,9 +51,6 @@ def getitem_next(self, wheretpe, isadvanced): if len(wheretpe.types) == 0: return self headtpe = wheretpe.types[0] - - print("getitem_next", headtpe) - tailtpe = numba.types.Tuple(wheretpe.types[1:]) if isinstance(headtpe, numba.types.Integer): return self.contenttpe.carry().getitem_next(tailtpe, isadvanced) @@ -212,13 +209,9 @@ def lower_getitem_tuple(context, builder, sig, args): rettpe, (arraytpe, wheretpe) = sig.return_type, sig.args arrayval, whereval = args - print("lower_getitem_tuple BEFORE", wheretpe) - wheretpe, whereval = util.preprocess_slicetuple(context, builder, wheretpe, whereval) nexttpe, nextval = util.wrap_for_slicetuple(context, builder, arraytpe, arrayval) - print("lower_getitem_tuple AFTER", wheretpe) - outtpe = nexttpe.getitem_next(wheretpe, False) outval = nexttpe.lower_getitem_next(context, builder, nexttpe, wheretpe, nextval, whereval, None) @@ -269,8 +262,6 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, util.debug(context, builder, arraytpe.startstpe, proxyin.starts) util.debug(context, builder, arraytpe.stopstpe, proxyin.stops) - print("headtpe", headtpe) - nextcarry = util.newindex64(context, builder, numba.int64, lenstarts) util.call(context, builder, kernel, (util.arrayptr(context, builder, util.index64tpe, nextcarry), diff --git a/awkward1/_numba/util.py b/awkward1/_numba/util.py index 4e0ec5edf1..8340e5c65d 100644 --- a/awkward1/_numba/util.py +++ b/awkward1/_numba/util.py @@ -35,7 +35,7 @@ def cast(context, builder, fromtpe, totpe, val): else: raise AssertionError("unrecognized bitwidth") if fromtpe.bitwidth < totpe.bitwidth: - return builder.zext(val, context.get_value_type(totpe)) + return builder.sext(val, context.get_value_type(totpe)) elif fromtpe.bitwidth > totpe.bitwidth: return builder.trunc(val, context.get_value_type(totpe)) else: From 8436f83fa85601ff577f00521bf7d023d9b90540 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 9 Oct 2019 21:16:57 -0500 Subject: [PATCH 30/36] Cleaned up all debugging code. --- awkward1/_numba/listarray.py | 7 +------ awkward1/_numba/listoffsetarray.py | 8 -------- awkward1/_numba/util.py | 12 ------------ tests/test_PR015_finish_up_getitem_in_numba.py | 16 ---------------- 4 files changed, 1 insertion(+), 42 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index 5934297f45..c1136c99d2 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -257,11 +257,6 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, else: raise AssertionError("unrecognized bitwidth") - tmp = util.cast(context, builder, headtpe, numba.int64, headval) - util.debug(context, builder, headtpe, headval, numba.int64, tmp) - util.debug(context, builder, arraytpe.startstpe, proxyin.starts) - util.debug(context, builder, arraytpe.stopstpe, proxyin.stops) - nextcarry = util.newindex64(context, builder, numba.int64, lenstarts) util.call(context, builder, kernel, (util.arrayptr(context, builder, util.index64tpe, nextcarry), @@ -270,7 +265,7 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, lenstarts, context.get_constant(numba.int64, 0), context.get_constant(numba.int64, 0), - tmp), + util.cast(context, builder, headtpe, numba.int64, headval)), "in {}, indexing error".format(arraytpe.shortname)) nextcontenttpe = arraytpe.contenttpe.carry() nextcontentval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index e5887b7701..1779d2ebed 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -269,9 +269,6 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, else: raise AssertionError("unrecognized bitwidth") - util.debug(context, builder, headtpe, headval) - util.debug(context, builder, arraytpe.offsetstpe, proxyin.offsets) - nextcarry = util.newindex64(context, builder, numba.int64, lenstarts) util.call(context, builder, kernel, (util.arrayptr(context, builder, util.index64tpe, nextcarry), @@ -350,11 +347,6 @@ def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): starts, stops = starts_stops(context, builder, arraytpe.offsetstpe, proxyin.offsets, lenstarts, lenoffsets) - util.debug(context, builder, numba.int64, lenoffsets, numba.int64, lenstarts) - util.debug(context, builder, arraytpe.offsetstpe, proxyin.offsets) - util.debug(context, builder, arraytpe.offsetstpe, starts) - util.debug(context, builder, arraytpe.offsetstpe, stops) - proxyout = numba.cgutils.create_struct_proxy(awkward1._numba.listarray.ListArrayType(arraytpe.offsetstpe, arraytpe.offsetstpe, arraytpe.contenttpe, arraytpe.idtpe))(context, builder) proxyout.starts = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (starts, carryval)) proxyout.stops = numba.targets.arrayobj.fancy_getitem_array(context, builder, arraytpe.offsetstpe(arraytpe.offsetstpe, carrytpe), (stops, carryval)) diff --git a/awkward1/_numba/util.py b/awkward1/_numba/util.py index 8340e5c65d..43652b7c73 100644 --- a/awkward1/_numba/util.py +++ b/awkward1/_numba/util.py @@ -141,15 +141,12 @@ def typing_broadcast_arrays(arrays): @numba.generated_jit(nopython=True) def regularize_slice(arrays): if not isinstance(arrays, numba.types.BaseTuple) and isinstance(arrays, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(arrays.dtype, numba.types.Boolean): - print("regularize_slice A", arrays) return lambda arrays: numpy.nonzero(arrays) elif not isinstance(arrays, numba.types.BaseTuple) or not any(isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) for t in arrays.types): - print("regularize_slice B", arrays) return lambda arrays: arrays else: - print("regularize_slice C", arrays) code = "def impl(arrays):\n" indexes = [] for i, t in enumerate(arrays.types): @@ -170,15 +167,12 @@ def regularize_slice(arrays): def typing_regularize_slice(arrays): out = () if not isinstance(arrays, numba.types.BaseTuple) and isinstance(arrays, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(arrays.dtype, numba.types.Boolean): - print("typing_regularize_slice A", arrays) return numba.types.Tuple(arrays.ndims*(numba.types.Array(numba.int64, 1, "C"),)) elif not isinstance(arrays, numba.types.BaseTuple) or not any(isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) for t in arrays.types): - print("typing_regularize_slice B", arrays) return arrays else: - print("typing_regularize_slice C", arrays) for t in arrays.types: if isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(t.dtype, numba.types.Boolean): out = out + t.ndims*(numba.types.Array(numba.int64, 1, "C"),) @@ -191,22 +185,16 @@ def typing_regularize_slice(arrays): return numba.types.Tuple(out) def preprocess_slicetuple(context, builder, wheretpe, whereval): - print("wheretpe1", wheretpe) - wheretpe2 = typing_regularize_slice(wheretpe) regularize_slice.compile(wheretpe2(wheretpe)) cres = regularize_slice.overloads[(wheretpe,)] whereval2 = context.call_internal(builder, cres.fndesc, wheretpe2(wheretpe), (whereval,)) - print("wheretpe2", wheretpe2) - wheretpe3 = typing_broadcast_arrays(wheretpe2) broadcast_arrays.compile(wheretpe3(wheretpe2)) cres2 = broadcast_arrays.overloads[(wheretpe2,)] whereval3 = context.call_internal(builder, cres2.fndesc, wheretpe3(wheretpe2), (whereval2,)) - print("wheretpe3", wheretpe3) - return wheretpe3, whereval3 def wrap_for_slicetuple(context, builder, arraytpe, arrayval): diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index 88a1592cb6..3669645750 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -121,16 +121,12 @@ def test_deep_listarray32(): listarrayB32 = awkward1.layout.ListArray32(awkward1.layout.Index32(startsB), awkward1.layout.Index32(stopsB), listarrayA32) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) - print("============================================") - @numba.njit def f1(q): return q[1, 2, 4] assert awkward1.tolist(f1(listarrayB32)) == awkward1.tolist(f1(modelB)) - print("--------------------------------------------") - @numba.njit def f2(q): return q[1, -1, 4, -2] @@ -148,16 +144,12 @@ def test_deep_listoffsetarray32(): listoffsetarrayB32 = awkward1.layout.ListOffsetArray32(awkward1.layout.Index32(offsetsB), listoffsetarrayA32) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) - print("============================================") - @numba.njit def f1(q): return q[1, 2, 4] assert awkward1.tolist(f1(listoffsetarrayB32)) == awkward1.tolist(f1(modelB)) - print("--------------------------------------------") - @numba.njit def f2(q): return q[1, -1, 4, -2] @@ -177,16 +169,12 @@ def test_deep_listarray64(): listarrayB64 = awkward1.layout.ListArray64(awkward1.layout.Index64(startsB), awkward1.layout.Index64(stopsB), listarrayA64) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) - print("============================================") - @numba.njit def f1(q): return q[1, 2, 4] assert awkward1.tolist(f1(listarrayB64)) == awkward1.tolist(f1(modelB)) - print("--------------------------------------------") - @numba.njit def f2(q): return q[1, -1, 4, -2] @@ -204,16 +192,12 @@ def test_deep_listoffsetarray64(): listoffsetarrayB64 = awkward1.layout.ListOffsetArray64(awkward1.layout.Index64(offsetsB), listoffsetarrayA64) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) - print("============================================") - @numba.njit def f1(q): return q[1, 2, 4] assert awkward1.tolist(f1(listoffsetarrayB64)) == awkward1.tolist(f1(modelB)) - print("--------------------------------------------") - @numba.njit def f2(q): return q[1, -1, 4, -2] From 35ee6c983339e6a9f479994a21672959dbf843ac Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Thu, 10 Oct 2019 06:45:10 -0500 Subject: [PATCH 31/36] [skip ci] in progress. --- awkward1/_numba/listarray.py | 38 +++++++++++++++++++ .../test_PR015_finish_up_getitem_in_numba.py | 24 ++++++++---- 2 files changed, 54 insertions(+), 8 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index c1136c99d2..14ec90b52e 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -272,6 +272,44 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, return nextcontenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, advanced) elif isinstance(headtpe, numba.types.SliceType): + proxyslicein = numba.cgutils.create_struct_proxy(headtpe)(context, builder, value=headval) + + if arraytpe.bitwidth == 64: + determine_carrylength = cpu.kernels.awkward_listarray32_getitem_next_range_carrylength + fill_carry = cpu.kernels.awkward_listarray32_getitem_next_range_64 + determine_total = cpu.kernels.awkward_listarray32_getitem_next_range_counts_64 + fill_nextadvanced = cpu.kernels.awkward_listarray32_getitem_next_range_spreadadvanced_64 + elif arraytpe.bitwidth == 32: + determine_carrylength = cpu.kernels.awkward_listarray32_getitem_next_range_carrylength + fill_carry = cpu.kernels.awkward_listarray32_getitem_next_range_64 + determine_total = cpu.kernels.awkward_listarray32_getitem_next_range_counts_64 + fill_nextadvanced = cpu.kernels.awkward_listarray32_getitem_next_range_spreadadvanced_64 + else: + raise AssertionError("unrecognized bitwidth") + + carrylength = numba.cgutils.alloca_once(builder, context.get_value_type(numba.int64)) + builder.store(context.get_constant(numba.int64, 123), carrylength) + util.debug(context, builder, numba.int64, builder.load(carrylength)) + + util.call(context, builder, determine_carrylength, + (carrylength, + util.arrayptr(context, builder, arraytpe.startstpe, proxyin.starts), + util.arrayptr(context, builder, arraytpe.stopstpe, proxyin.stops), + lenstarts, + context.get_constant(numba.int64, 0), + context.get_constant(numba.int64, 0), + util.cast(context, builder, numba.intp, numba.int64, proxyin.start), + util.cast(context, builder, numba.intp, numba.int64, proxyin.stop), + util.cast(context, builder, numba.intp, numba.int64, proxyin.step)), + "in {}, indexing error".format(arraytpe.shortname)) + + HERE HERE HERE HERE HERE HERE HERE HERE + + proxyout = numba.cgutils.create_struct_proxy(awkward1._numba.listoffsetarray.ListOffsetArrayType(arraytpe.startstpe, arraytpe.contenttpe, arraytpe.idtpe))(context, builder) + proxyout.offsets = proxyin.starts + proxyout.content = proxyin.content + return proxyout._getvalue() + raise NotImplementedError("ListArray.getitem_next(slice)") elif isinstance(headtpe, numba.types.EllipsisType): diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index 3669645750..b2ed6ba524 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -169,17 +169,25 @@ def test_deep_listarray64(): listarrayB64 = awkward1.layout.ListArray64(awkward1.layout.Index64(startsB), awkward1.layout.Index64(stopsB), listarrayA64) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) - @numba.njit - def f1(q): - return q[1, 2, 4] - - assert awkward1.tolist(f1(listarrayB64)) == awkward1.tolist(f1(modelB)) + # @numba.njit + # def f1(q): + # return q[1, 2, 4] + # + # assert awkward1.tolist(f1(listarrayB64)) == awkward1.tolist(f1(modelB)) + # + # @numba.njit + # def f2(q): + # return q[1, -1, 4, -2] + # + # assert f2(listarrayB64) == f2(modelB) @numba.njit - def f2(q): - return q[1, -1, 4, -2] + def f3(q): + # return q[1:, ::2, 1:-1, 5::-1] + return q[1:,] - assert f2(listarrayB64) == f2(modelB) + print(awkward1.tolist(f3(listarrayB64))) + raise Exception def test_deep_listoffsetarray64(): content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) From 37e277b6c5be4ba5225aafb92980f2983d8476cf Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Thu, 10 Oct 2019 08:30:53 -0500 Subject: [PATCH 32/36] Partially implemented ListArray.getitem_next(range). --- awkward1/_numba/listarray.py | 78 +++++++++++++++---- awkward1/_numba/listoffsetarray.py | 21 +++-- .../test_PR015_finish_up_getitem_in_numba.py | 28 ++++--- 3 files changed, 91 insertions(+), 36 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index 14ec90b52e..aa2d5496f2 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -43,8 +43,8 @@ def getitem_range(self): def getitem_tuple(self, wheretpe): nexttpe = ListArrayType(util.index64tpe, util.index64tpe, self, numba.none) - out = nexttpe.getitem_next(wheretpe, False) - return out.getitem_int() + outtpe = nexttpe.getitem_next(wheretpe, False) + return outtpe.getitem_int() def getitem_next(self, wheretpe, isadvanced): import awkward1._numba.listoffsetarray @@ -52,21 +52,32 @@ def getitem_next(self, wheretpe, isadvanced): return self headtpe = wheretpe.types[0] tailtpe = numba.types.Tuple(wheretpe.types[1:]) + if isinstance(headtpe, numba.types.Integer): return self.contenttpe.carry().getitem_next(tailtpe, isadvanced) + elif isinstance(headtpe, numba.types.SliceType): - return awkward1._numba.listoffsetarray.ListOffsetArrayType(self.startstpe, self.contenttpe.getitem_next(tailtpe, isadvanced), self.idtpe) + contenttpe = self.contenttpe.carry().getitem_next(tailtpe, isadvanced) + if not isadvanced: + return awkward1._numba.listoffsetarray.ListOffsetArrayType(util.index64tpe, contenttpe, self.idtpe) + else: + raise NotImplementedError("ListArray isadvanced") + elif isinstance(headtpe, numba.types.EllipsisType): raise NotImplementedError("ellipsis") + elif isinstance(headtpe, type(numba.typeof(numpy.newaxis))): raise NotImplementedError("newaxis") - elif isinstance(headtpe, numba.types.Array) and not isadvanced: + + elif isinstance(headtpe, numba.types.Array): if headtpe.ndim != 1: raise NotImplementedError("array.ndim != 1") contenttpe = self.contenttpe.carry().getitem_next(tailtpe, True) - return awkward1._numba.listoffsetarray.ListOffsetArrayType(self.startstpe, contenttpe, self.idtpe) - elif isinstance(headtpe, numba.types.Array): - return self.contenttpe.getitem_next(tailtpe, True) + if not isadvanced: + return awkward1._numba.listoffsetarray.ListOffsetArrayType(util.index64tpe, contenttpe, self.idtpe) + else: + return contenttpe + else: raise AssertionError(headtpe) @@ -275,10 +286,10 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, proxyslicein = numba.cgutils.create_struct_proxy(headtpe)(context, builder, value=headval) if arraytpe.bitwidth == 64: - determine_carrylength = cpu.kernels.awkward_listarray32_getitem_next_range_carrylength - fill_carry = cpu.kernels.awkward_listarray32_getitem_next_range_64 - determine_total = cpu.kernels.awkward_listarray32_getitem_next_range_counts_64 - fill_nextadvanced = cpu.kernels.awkward_listarray32_getitem_next_range_spreadadvanced_64 + determine_carrylength = cpu.kernels.awkward_listarray64_getitem_next_range_carrylength + fill_carry = cpu.kernels.awkward_listarray64_getitem_next_range_64 + determine_total = cpu.kernels.awkward_listarray64_getitem_next_range_counts_64 + fill_nextadvanced = cpu.kernels.awkward_listarray64_getitem_next_range_spreadadvanced_64 elif arraytpe.bitwidth == 32: determine_carrylength = cpu.kernels.awkward_listarray32_getitem_next_range_carrylength fill_carry = cpu.kernels.awkward_listarray32_getitem_next_range_64 @@ -288,6 +299,7 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, raise AssertionError("unrecognized bitwidth") carrylength = numba.cgutils.alloca_once(builder, context.get_value_type(numba.int64)) + builder.store(context.get_constant(numba.int64, 123), carrylength) util.debug(context, builder, numba.int64, builder.load(carrylength)) @@ -298,12 +310,45 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, lenstarts, context.get_constant(numba.int64, 0), context.get_constant(numba.int64, 0), - util.cast(context, builder, numba.intp, numba.int64, proxyin.start), - util.cast(context, builder, numba.intp, numba.int64, proxyin.stop), - util.cast(context, builder, numba.intp, numba.int64, proxyin.step)), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.start), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.stop), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.step)), "in {}, indexing error".format(arraytpe.shortname)) - HERE HERE HERE HERE HERE HERE HERE HERE + util.debug(context, builder, numba.int64, builder.load(carrylength)) + + nextoffsets = util.newindex64(context, builder, numba.int64, builder.add(lenstarts, context.get_constant(numba.int64, 1))) + nextcarry = util.newindex64(context, builder, numba.int64, builder.load(carrylength)) + + util.call(context, builder, fill_carry, + (util.arrayptr(context, builder, util.index64tpe, nextoffsets), + util.arrayptr(context, builder, util.index64tpe, nextcarry), + util.arrayptr(context, builder, arraytpe.startstpe, proxyin.starts), + util.arrayptr(context, builder, arraytpe.stopstpe, proxyin.stops), + lenstarts, + context.get_constant(numba.int64, 0), + context.get_constant(numba.int64, 0), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.start), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.stop), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.step)), + "in {}, indexing error".format(arraytpe.shortname)) + + nextcontenttpe = arraytpe.contenttpe.carry() + nextcontentval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) + + if advanced is None: + outcontenttpe = nextcontenttpe.getitem_next(tailtpe, False) + outcontentval = nextcontenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, advanced) + outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(util.index64tpe, outcontenttpe, arraytpe.idtpe) + proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) + proxyout.offsets = nextoffsets + proxyout.content = outcontentval + if arraytpe.idtpe != numba.none: + proxyout.id = proxyin.id + return proxyout._getvalue() + + else: + raise Exception proxyout = numba.cgutils.create_struct_proxy(awkward1._numba.listoffsetarray.ListOffsetArrayType(arraytpe.startstpe, arraytpe.contenttpe, arraytpe.idtpe))(context, builder) proxyout.offsets = proxyin.starts @@ -359,7 +404,7 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, contenttpe = nexttpe.getitem_next(tailtpe, True) contentval = nexttpe.lower_getitem_next(context, builder, nexttpe, tailtpe, nextval, tailval, nextadvanced) - outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(arraytpe.startstpe, contenttpe, arraytpe.idtpe) + outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(util.index64tpe, contenttpe, arraytpe.idtpe) proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) proxyout.offsets = nextoffsets proxyout.content = contentval @@ -411,6 +456,7 @@ def lower_carry(context, builder, arraytpe, carrytpe, arrayval, carryval): proxyout.content = proxyin.content if arraytpe.idtpe != numba.none: proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, arraytpe.idtpe, carrytpe, proxyin.id, carryval) + return proxyout._getvalue() @numba.typing.templates.infer_getattr diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index 1779d2ebed..71870b5ebf 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -46,21 +46,32 @@ def getitem_next(self, wheretpe, isadvanced): return self headtpe = wheretpe.types[0] tailtpe = numba.types.Tuple(wheretpe.types[1:]) + if isinstance(headtpe, numba.types.Integer): return self.contenttpe.carry().getitem_next(tailtpe, isadvanced) + elif isinstance(headtpe, numba.types.SliceType): - return ListOffsetArrayType(self.offsetstpe, self.contenttpe.getitem_next(tailtpe, isadvanced), self.idtpe) + contenttpe = self.contenttpe.carry().getitem_next(tailtpe, isadvanced) + if not isadvanced: + return ListOffsetArrayType(util.index64tpe, contenttpe, self.idtpe) + else: + raise NotImplementedError("ListOffsetArray isadvanced") + elif isinstance(headtpe, numba.types.EllipsisType): raise NotImplementedError("ellipsis") + elif isinstance(headtpe, type(numba.typeof(numpy.newaxis))): raise NotImplementedError("newaxis") - elif isinstance(headtpe, numba.types.Array) and not advanced: + + elif isinstance(headtpe, numba.types.Array): if headtpe.ndim != 1: raise NotImplementedError("array.ndim != 1") contenttpe = self.contenttpe.carry().getitem_next(tailtpe, True) - return ListOffsetArrayType(self.startstpe, contenttpe, self.idtpe) - elif isinstance(headtpe, numba.types.Array): - return self.contenttpe.getitem_next(tailtpe, True) + if not isadvanced: + return ListOffsetArrayType(util.index64tpe, contenttpe, self.idtpe) + else: + return contenttpe + else: raise AssertionError(headtpe) diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index b2ed6ba524..21d6e54322 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -169,25 +169,23 @@ def test_deep_listarray64(): listarrayB64 = awkward1.layout.ListArray64(awkward1.layout.Index64(startsB), awkward1.layout.Index64(stopsB), listarrayA64) modelB = numpy.arange(2*3*5*7).reshape(2, 3, 5, 7) - # @numba.njit - # def f1(q): - # return q[1, 2, 4] - # - # assert awkward1.tolist(f1(listarrayB64)) == awkward1.tolist(f1(modelB)) - # - # @numba.njit - # def f2(q): - # return q[1, -1, 4, -2] - # - # assert f2(listarrayB64) == f2(modelB) + @numba.njit + def f1(q): + return q[1, 2, 4] + + assert awkward1.tolist(f1(listarrayB64)) == awkward1.tolist(f1(modelB)) + + @numba.njit + def f2(q): + return q[1, -1, 4, -2] + + assert f2(listarrayB64) == f2(modelB) @numba.njit def f3(q): - # return q[1:, ::2, 1:-1, 5::-1] - return q[1:,] + return q[1:, ::2, 1:-1, 5::-1] - print(awkward1.tolist(f3(listarrayB64))) - raise Exception + assert awkward1.tolist(f3(listarrayB64)) == awkward1.tolist(f3(modelB)) def test_deep_listoffsetarray64(): content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) From 7552692d9d006038d06bfc0f4fa68ba133e8239d Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Thu, 10 Oct 2019 09:02:52 -0500 Subject: [PATCH 33/36] Homogenized ListArray and ListOffsetArray; ensured that nextoffsets follows bit width of old starts or offsets. --- awkward1/_numba/listarray.py | 31 ++--- awkward1/_numba/listoffsetarray.py | 109 ++++++++++++------ awkward1/_numba/util.py | 24 ++++ .../test_PR015_finish_up_getitem_in_numba.py | 6 + 4 files changed, 114 insertions(+), 56 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index aa2d5496f2..c87b26193c 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -59,7 +59,7 @@ def getitem_next(self, wheretpe, isadvanced): elif isinstance(headtpe, numba.types.SliceType): contenttpe = self.contenttpe.carry().getitem_next(tailtpe, isadvanced) if not isadvanced: - return awkward1._numba.listoffsetarray.ListOffsetArrayType(util.index64tpe, contenttpe, self.idtpe) + return awkward1._numba.listoffsetarray.ListOffsetArrayType(util.indextpe(self.bitwidth), contenttpe, self.idtpe) else: raise NotImplementedError("ListArray isadvanced") @@ -74,7 +74,7 @@ def getitem_next(self, wheretpe, isadvanced): raise NotImplementedError("array.ndim != 1") contenttpe = self.contenttpe.carry().getitem_next(tailtpe, True) if not isadvanced: - return awkward1._numba.listoffsetarray.ListOffsetArrayType(util.index64tpe, contenttpe, self.idtpe) + return awkward1._numba.listoffsetarray.ListOffsetArrayType(util.indextpe(self.bitwidth), contenttpe, self.idtpe) else: return contenttpe @@ -299,10 +299,6 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, raise AssertionError("unrecognized bitwidth") carrylength = numba.cgutils.alloca_once(builder, context.get_value_type(numba.int64)) - - builder.store(context.get_constant(numba.int64, 123), carrylength) - util.debug(context, builder, numba.int64, builder.load(carrylength)) - util.call(context, builder, determine_carrylength, (carrylength, util.arrayptr(context, builder, arraytpe.startstpe, proxyin.starts), @@ -315,13 +311,10 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, util.cast(context, builder, numba.intp, numba.int64, proxyslicein.step)), "in {}, indexing error".format(arraytpe.shortname)) - util.debug(context, builder, numba.int64, builder.load(carrylength)) - - nextoffsets = util.newindex64(context, builder, numba.int64, builder.add(lenstarts, context.get_constant(numba.int64, 1))) + nextoffsets = util.newindex(arraytpe.bitwidth, context, builder, numba.int64, builder.add(lenstarts, context.get_constant(numba.int64, 1))) nextcarry = util.newindex64(context, builder, numba.int64, builder.load(carrylength)) - util.call(context, builder, fill_carry, - (util.arrayptr(context, builder, util.index64tpe, nextoffsets), + (util.arrayptr(context, builder, util.indextpe(arraytpe.bitwidth), nextoffsets), util.arrayptr(context, builder, util.index64tpe, nextcarry), util.arrayptr(context, builder, arraytpe.startstpe, proxyin.starts), util.arrayptr(context, builder, arraytpe.stopstpe, proxyin.stops), @@ -339,7 +332,7 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, if advanced is None: outcontenttpe = nextcontenttpe.getitem_next(tailtpe, False) outcontentval = nextcontenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, advanced) - outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(util.index64tpe, outcontenttpe, arraytpe.idtpe) + outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(util.indextpe(arraytpe.bitwidth), outcontenttpe, arraytpe.idtpe) proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) proxyout.offsets = nextoffsets proxyout.content = outcontentval @@ -350,10 +343,10 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, else: raise Exception - proxyout = numba.cgutils.create_struct_proxy(awkward1._numba.listoffsetarray.ListOffsetArrayType(arraytpe.startstpe, arraytpe.contenttpe, arraytpe.idtpe))(context, builder) - proxyout.offsets = proxyin.starts - proxyout.content = proxyin.content - return proxyout._getvalue() + # proxyout = numba.cgutils.create_struct_proxy(awkward1._numba.listoffsetarray.ListOffsetArrayType(arraytpe.startstpe, arraytpe.contenttpe, arraytpe.idtpe))(context, builder) + # proxyout.offsets = proxyin.starts + # proxyout.content = proxyin.content + # return proxyout._getvalue() raise NotImplementedError("ListArray.getitem_next(slice)") @@ -383,9 +376,9 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, nextcarry = util.newindex64(context, builder, numba.int64, lencarry) nextadvanced = util.newindex64(context, builder, numba.int64, lencarry) - nextoffsets = util.newindex64(context, builder, numba.int64, lenoffsets) + nextoffsets = util.newindex(arraytpe.bitwidth, context, builder, numba.int64, lenoffsets) util.call(context, builder, kernel, - (util.arrayptr(context, builder, util.index64tpe, nextoffsets), + (util.arrayptr(context, builder, util.indextpe(arraytpe.bitwidth), nextoffsets), util.arrayptr(context, builder, util.index64tpe, nextcarry), util.arrayptr(context, builder, util.index64tpe, nextadvanced), util.arrayptr(context, builder, arraytpe.startstpe, proxyin.starts), @@ -404,7 +397,7 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, contenttpe = nexttpe.getitem_next(tailtpe, True) contentval = nexttpe.lower_getitem_next(context, builder, nexttpe, tailtpe, nextval, tailval, nextadvanced) - outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(util.index64tpe, contenttpe, arraytpe.idtpe) + outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(util.indextpe(arraytpe.bitwidth), contenttpe, arraytpe.idtpe) proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) proxyout.offsets = nextoffsets proxyout.content = contentval diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index 71870b5ebf..ace4ed0fec 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -53,7 +53,7 @@ def getitem_next(self, wheretpe, isadvanced): elif isinstance(headtpe, numba.types.SliceType): contenttpe = self.contenttpe.carry().getitem_next(tailtpe, isadvanced) if not isadvanced: - return ListOffsetArrayType(util.index64tpe, contenttpe, self.idtpe) + return ListOffsetArrayType(util.indextpe(self.bitwidth), contenttpe, self.idtpe) else: raise NotImplementedError("ListOffsetArray isadvanced") @@ -68,7 +68,7 @@ def getitem_next(self, wheretpe, isadvanced): raise NotImplementedError("array.ndim != 1") contenttpe = self.contenttpe.carry().getitem_next(tailtpe, True) if not isadvanced: - return ListOffsetArrayType(util.index64tpe, contenttpe, self.idtpe) + return ListOffsetArrayType(util.indextpe(self.bitwidth), contenttpe, self.idtpe) else: return contenttpe @@ -303,48 +303,83 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, elif isinstance(headtpe, type(numba.typeof(numpy.newaxis))): raise NotImplementedError("ListOffsetArray.getitem_next(newaxis)") - elif isinstance(headtpe, numba.types.Array) and advanced is None: + elif isinstance(headtpe, numba.types.Array): if headtpe.ndim != 1: raise NotImplementedError("array.ndim != 1") flathead = numba.targets.arrayobj.array_ravel(context, builder, util.int64tep(headtpe), (headval,)) lenflathead = util.arraylen(context, builder, util.int64tpe, flathead, totpe=numba.int64) - lencarry = builder.mul(lenstarts, lenflathead) - - nextcarry = util.newindex64(context, builder, numba.int64, lencarry) - nextadvanced = util.newindex64(context, builder, numba.int64, lencarry) - nextoffsets = util.newindex64(context, builder, numba.int64, lenoffsets) - - util.call(context, builder, cpu.kernels.awkward_listarray64_getitem_next_array_64, - (util.arrayptr(context, builder, util.index64tpe, nextoffsets), - util.arrayptr(context, builder, util.index64tpe, nextcarry), - util.arrayptr(context, builder, util.index64tpe, nextadvanced), - util.arrayptr(context, builder, arraytpe.offsetstpe, starts), - util.arrayptr(context, builder, arraytpe.offsetstpe, stops), - util.arrayptr(context, builder, util.index64tpe, flathead), - context.get_constant(numba.int64, 0), - context.get_constant(numba.int64, 0), - lenstarts, - lenflathead, - lencontent), - "in {}, indexing error".format(arraytpe.shortname)) - - nexttpe = arraytpe.contenttpe.carry() - nextval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) - contenttpe = nexttpe.getitem_next(tailtpe, True) - contentval = nexttpe.lower_getitem_next(context, builder, nexttpe, tailtpe, nextval, tailval, nextadvanced) - - outtpe = ListOffsetArrayType(arraytpe.offsetstpe, contenttpe, arraytpe.idtpe) - proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) - proxyout.offsets = nextoffsets - proxyout.content = contentval - if outtpe.idtpe != numba.none: - proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, outtpe.idtpe, util.index64tpe, proxyin.id, flathead) - return proxyout._getvalue() + if advanced is None: + if arraytpe.bitwidth == 64: + kernel = cpu.kernels.awkward_listarray64_getitem_next_array_64 + elif arraytpe.bitwidth == 32: + kernel = cpu.kernels.awkward_listarray32_getitem_next_array_64 + else: + raise AssertionError("unrecognized bitwidth") + + lencarry = builder.mul(lenstarts, lenflathead) + + nextcarry = util.newindex64(context, builder, numba.int64, lencarry) + nextadvanced = util.newindex64(context, builder, numba.int64, lencarry) + nextoffsets = util.newindex(arraytpe.bitwidth, context, builder, numba.int64, lenoffsets) + util.call(context, builder, kernel, + (util.arrayptr(context, builder, util.indextpe(arraytpe.bitwidth), nextoffsets), + util.arrayptr(context, builder, util.index64tpe, nextcarry), + util.arrayptr(context, builder, util.index64tpe, nextadvanced), + util.arrayptr(context, builder, arraytpe.offsetstpe, starts), + util.arrayptr(context, builder, arraytpe.offsetstpe, stops), + util.arrayptr(context, builder, util.index64tpe, flathead), + context.get_constant(numba.int64, 0), + context.get_constant(numba.int64, 0), + lenstarts, + lenflathead, + lencontent), + "in {}, indexing error".format(arraytpe.shortname)) + + nexttpe = arraytpe.contenttpe.carry() + nextval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) + + contenttpe = nexttpe.getitem_next(tailtpe, True) + contentval = nexttpe.lower_getitem_next(context, builder, nexttpe, tailtpe, nextval, tailval, nextadvanced) + + outtpe = ListOffsetArrayType(util.indextpe(arraytpe.bitwidth), contenttpe, arraytpe.idtpe) + proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) + proxyout.offsets = nextoffsets + proxyout.content = contentval + if outtpe.idtpe != numba.none: + proxyout.id = awkward1._numba.identity.lower_getitem_any(context, builder, outtpe.idtpe, util.index64tpe, proxyin.id, flathead) + return proxyout._getvalue() - elif isinstance(headtpe, numba.types.Array): - raise NotImplementedError("ListArray.getitem_next(advanced Array)") + else: + if arraytpe.bitwidth == 64: + kernel = cpu.kernels.awkward_listarray64_getitem_next_array_advanced_64 + elif arraytpe.bitwidth == 32: + kernel = cpu.kernels.awkward_listarray32_getitem_next_array_advanced_64 + else: + raise AssertionError("unrecognized bitwidth") + + nextcarry = util.newindex64(context, builder, numba.int64, lenstarts) + nextadvanced = util.newindex64(context, builder, numba.int64, lenstarts) + util.call(context, builder, kernel, + (util.arrayptr(context, builder, util.index64tpe, nextcarry), + util.arrayptr(context, builder, util.index64tpe, nextadvanced), + util.arrayptr(context, builder, arraytpe.offsetstpe, starts), + util.arrayptr(context, builder, arraytpe.offsetstpe, stops), + util.arrayptr(context, builder, util.index64tpe, flathead), + util.arrayptr(context, builder, util.index64tpe, advanced), + context.get_constant(numba.int64, 0), + context.get_constant(numba.int64, 0), + lenstarts, + lenflathead, + lencontent), + "in {}, indexing error".format(arraytpe.shortname)) + + nexttpe = arraytpe.contenttpe.carry() + nextval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) + + outtpe = nexttpe.getitem_next(tailtpe, True) + return nexttpe.lower_getitem_next(context, builder, nexttpe, tailtpe, nextval, tailval, nextadvanced) else: raise AssertionError(headtpe) diff --git a/awkward1/_numba/util.py b/awkward1/_numba/util.py index 43652b7c73..40b10a225f 100644 --- a/awkward1/_numba/util.py +++ b/awkward1/_numba/util.py @@ -12,7 +12,18 @@ RefType = numba.int64 +index8tpe = numba.types.Array(numba.int8, 1, "C") +index32tpe = numba.types.Array(numba.int32, 1, "C") index64tpe = numba.types.Array(numba.int64, 1, "C") +def indextpe(bitwidth): + if bitwidth == 8: + return index8tpe + elif bitwidth == 32: + return index32tpe + elif bitwidth == 64: + return index64tpe + else: + raise AssertionError(bitwidth) if not py27: exec(""" @@ -79,8 +90,21 @@ def call(context, builder, fcn, args, errormessage=None): # else: # raise AssertionError("unrecognized exception calling convention: {}".format(excptr)) +def newindex8(context, builder, lentpe, lenval): + return numba.targets.arrayobj.numpy_empty_nd(context, builder, index8tpe(lentpe), (lenval,)) +def newindex32(context, builder, lentpe, lenval): + return numba.targets.arrayobj.numpy_empty_nd(context, builder, index32tpe(lentpe), (lenval,)) def newindex64(context, builder, lentpe, lenval): return numba.targets.arrayobj.numpy_empty_nd(context, builder, index64tpe(lentpe), (lenval,)) +def newindex(bitwidth, context, builder, lentpe, lenval): + if bitwidth == 8: + return newindex8(context, builder, lentpe, lenval) + elif bitwidth == 32: + return newindex32(context, builder, lentpe, lenval) + elif bitwidth == 64: + return newindex64(context, builder, lentpe, lenval) + else: + raise AssertionError(bitwidth) @numba.jit(nopython=True) def shapeat(shapeat, array, at, ndim): diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index 21d6e54322..dd5b6a5725 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -133,6 +133,12 @@ def f2(q): assert f2(listarrayB32) == f2(modelB) + @numba.njit + def f3(q): + return q[1:, ::2, 1:-1, 5::-1] + + assert awkward1.tolist(f3(listarrayB32)) == awkward1.tolist(f3(modelB)) + def test_deep_listoffsetarray32(): content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) From c6033bc9d293d1651fb34f969bf9a0ea0595b7df Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Thu, 10 Oct 2019 09:07:51 -0500 Subject: [PATCH 34/36] Implemented ListOffsetArray.getitem_next(range, not advanced). --- awkward1/_numba/listoffsetarray.py | 60 ++++++++++++++++++- .../test_PR015_finish_up_getitem_in_numba.py | 12 ++++ 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index ace4ed0fec..f7c7c8d3d7 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -295,7 +295,65 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, return nextcontenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, advanced) elif isinstance(headtpe, numba.types.SliceType): - raise NotImplementedError("ListOffsetArray.getitem_next(slice)") + proxyslicein = numba.cgutils.create_struct_proxy(headtpe)(context, builder, value=headval) + + if arraytpe.bitwidth == 64: + determine_carrylength = cpu.kernels.awkward_listarray64_getitem_next_range_carrylength + fill_carry = cpu.kernels.awkward_listarray64_getitem_next_range_64 + determine_total = cpu.kernels.awkward_listarray64_getitem_next_range_counts_64 + fill_nextadvanced = cpu.kernels.awkward_listarray64_getitem_next_range_spreadadvanced_64 + elif arraytpe.bitwidth == 32: + determine_carrylength = cpu.kernels.awkward_listarray32_getitem_next_range_carrylength + fill_carry = cpu.kernels.awkward_listarray32_getitem_next_range_64 + determine_total = cpu.kernels.awkward_listarray32_getitem_next_range_counts_64 + fill_nextadvanced = cpu.kernels.awkward_listarray32_getitem_next_range_spreadadvanced_64 + else: + raise AssertionError("unrecognized bitwidth") + + carrylength = numba.cgutils.alloca_once(builder, context.get_value_type(numba.int64)) + util.call(context, builder, determine_carrylength, + (carrylength, + util.arrayptr(context, builder, arraytpe.offsetstpe, starts), + util.arrayptr(context, builder, arraytpe.offsetstpe, stops), + lenstarts, + context.get_constant(numba.int64, 0), + context.get_constant(numba.int64, 0), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.start), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.stop), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.step)), + "in {}, indexing error".format(arraytpe.shortname)) + + nextoffsets = util.newindex(arraytpe.bitwidth, context, builder, numba.int64, builder.add(lenstarts, context.get_constant(numba.int64, 1))) + nextcarry = util.newindex64(context, builder, numba.int64, builder.load(carrylength)) + util.call(context, builder, fill_carry, + (util.arrayptr(context, builder, util.indextpe(arraytpe.bitwidth), nextoffsets), + util.arrayptr(context, builder, util.index64tpe, nextcarry), + util.arrayptr(context, builder, arraytpe.offsetstpe, starts), + util.arrayptr(context, builder, arraytpe.offsetstpe, stops), + lenstarts, + context.get_constant(numba.int64, 0), + context.get_constant(numba.int64, 0), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.start), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.stop), + util.cast(context, builder, numba.intp, numba.int64, proxyslicein.step)), + "in {}, indexing error".format(arraytpe.shortname)) + + nextcontenttpe = arraytpe.contenttpe.carry() + nextcontentval = arraytpe.contenttpe.lower_carry(context, builder, arraytpe.contenttpe, util.index64tpe, proxyin.content, nextcarry) + + if advanced is None: + outcontenttpe = nextcontenttpe.getitem_next(tailtpe, False) + outcontentval = nextcontenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, advanced) + outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(util.indextpe(arraytpe.bitwidth), outcontenttpe, arraytpe.idtpe) + proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) + proxyout.offsets = nextoffsets + proxyout.content = outcontentval + if arraytpe.idtpe != numba.none: + proxyout.id = proxyin.id + return proxyout._getvalue() + + else: + raise Exception elif isinstance(headtpe, numba.types.EllipsisType): raise NotImplementedError("ListOffsetArray.getitem_next(ellipsis)") diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index dd5b6a5725..8d9b2ed50d 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -162,6 +162,12 @@ def f2(q): assert f2(listoffsetarrayB32) == f2(modelB) + @numba.njit + def f3(q): + return q[1:, ::2, 1:-1, 5::-1] + + assert awkward1.tolist(f3(listoffsetarrayB32)) == awkward1.tolist(f3(modelB)) + def test_deep_listarray64(): content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) @@ -215,3 +221,9 @@ def f2(q): return q[1, -1, 4, -2] assert f2(listoffsetarrayB64) == f2(modelB) + + @numba.njit + def f3(q): + return q[1:, ::2, 1:-1, 5::-1] + + assert awkward1.tolist(f3(listoffsetarrayB64)) == awkward1.tolist(f3(modelB)) From df3afaef91fac8bef2087b382123dcb7c6a86e20 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Thu, 10 Oct 2019 09:44:30 -0500 Subject: [PATCH 35/36] All of the Numba getitems seem to be done. --- awkward1/_numba/listarray.py | 41 +++++++++++-------- awkward1/_numba/listoffsetarray.py | 38 +++++++++++------ awkward1/_numba/util.py | 4 +- .../test_PR015_finish_up_getitem_in_numba.py | 24 +++++++++++ 4 files changed, 76 insertions(+), 31 deletions(-) diff --git a/awkward1/_numba/listarray.py b/awkward1/_numba/listarray.py index c87b26193c..37ae7b81fa 100644 --- a/awkward1/_numba/listarray.py +++ b/awkward1/_numba/listarray.py @@ -58,10 +58,7 @@ def getitem_next(self, wheretpe, isadvanced): elif isinstance(headtpe, numba.types.SliceType): contenttpe = self.contenttpe.carry().getitem_next(tailtpe, isadvanced) - if not isadvanced: - return awkward1._numba.listoffsetarray.ListOffsetArrayType(util.indextpe(self.bitwidth), contenttpe, self.idtpe) - else: - raise NotImplementedError("ListArray isadvanced") + return awkward1._numba.listoffsetarray.ListOffsetArrayType(util.indextpe(self.bitwidth), contenttpe, self.idtpe) elif isinstance(headtpe, numba.types.EllipsisType): raise NotImplementedError("ellipsis") @@ -332,23 +329,33 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, if advanced is None: outcontenttpe = nextcontenttpe.getitem_next(tailtpe, False) outcontentval = nextcontenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, advanced) - outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(util.indextpe(arraytpe.bitwidth), outcontenttpe, arraytpe.idtpe) - proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) - proxyout.offsets = nextoffsets - proxyout.content = outcontentval - if arraytpe.idtpe != numba.none: - proxyout.id = proxyin.id - return proxyout._getvalue() else: - raise Exception + total = numba.cgutils.alloca_once(builder, context.get_value_type(numba.int64)) + util.call(context, builder, determine_total, + (total, + util.arrayptr(context, builder, util.indextpe(arraytpe.bitwidth), nextoffsets), + lenstarts), + "in {}, indexing error".format(arraytpe.shortname)) + + nextadvanced = util.newindex64(context, builder, numba.int64, builder.load(total)) + util.call(context, builder, fill_nextadvanced, + (util.arrayptr(context, builder, util.index64tpe, nextadvanced), + util.arrayptr(context, builder, util.index64tpe, advanced), + util.arrayptr(context, builder, util.indextpe(arraytpe.bitwidth), nextoffsets), + lenstarts), + "in {}, indexing error".format(arraytpe.shortname)) - # proxyout = numba.cgutils.create_struct_proxy(awkward1._numba.listoffsetarray.ListOffsetArrayType(arraytpe.startstpe, arraytpe.contenttpe, arraytpe.idtpe))(context, builder) - # proxyout.offsets = proxyin.starts - # proxyout.content = proxyin.content - # return proxyout._getvalue() + outcontenttpe = nextcontenttpe.getitem_next(tailtpe, True) + outcontentval = nextcontenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, nextadvanced) - raise NotImplementedError("ListArray.getitem_next(slice)") + outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(util.indextpe(arraytpe.bitwidth), outcontenttpe, arraytpe.idtpe) + proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) + proxyout.offsets = nextoffsets + proxyout.content = outcontentval + if arraytpe.idtpe != numba.none: + proxyout.id = proxyin.id + return proxyout._getvalue() elif isinstance(headtpe, numba.types.EllipsisType): raise NotImplementedError("ListArray.getitem_next(ellipsis)") diff --git a/awkward1/_numba/listoffsetarray.py b/awkward1/_numba/listoffsetarray.py index f7c7c8d3d7..e6a1e21cd5 100644 --- a/awkward1/_numba/listoffsetarray.py +++ b/awkward1/_numba/listoffsetarray.py @@ -52,10 +52,7 @@ def getitem_next(self, wheretpe, isadvanced): elif isinstance(headtpe, numba.types.SliceType): contenttpe = self.contenttpe.carry().getitem_next(tailtpe, isadvanced) - if not isadvanced: - return ListOffsetArrayType(util.indextpe(self.bitwidth), contenttpe, self.idtpe) - else: - raise NotImplementedError("ListOffsetArray isadvanced") + return ListOffsetArrayType(util.indextpe(self.bitwidth), contenttpe, self.idtpe) elif isinstance(headtpe, numba.types.EllipsisType): raise NotImplementedError("ellipsis") @@ -344,16 +341,33 @@ def lower_getitem_next(context, builder, arraytpe, wheretpe, arrayval, whereval, if advanced is None: outcontenttpe = nextcontenttpe.getitem_next(tailtpe, False) outcontentval = nextcontenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, advanced) - outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(util.indextpe(arraytpe.bitwidth), outcontenttpe, arraytpe.idtpe) - proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) - proxyout.offsets = nextoffsets - proxyout.content = outcontentval - if arraytpe.idtpe != numba.none: - proxyout.id = proxyin.id - return proxyout._getvalue() else: - raise Exception + total = numba.cgutils.alloca_once(builder, context.get_value_type(numba.int64)) + util.call(context, builder, determine_total, + (total, + util.arrayptr(context, builder, util.indextpe(arraytpe.bitwidth), nextoffsets), + lenstarts), + "in {}, indexing error".format(arraytpe.shortname)) + + nextadvanced = util.newindex64(context, builder, numba.int64, builder.load(total)) + util.call(context, builder, fill_nextadvanced, + (util.arrayptr(context, builder, util.index64tpe, nextadvanced), + util.arrayptr(context, builder, util.index64tpe, advanced), + util.arrayptr(context, builder, util.indextpe(arraytpe.bitwidth), nextoffsets), + lenstarts), + "in {}, indexing error".format(arraytpe.shortname)) + + outcontenttpe = nextcontenttpe.getitem_next(tailtpe, True) + outcontentval = nextcontenttpe.lower_getitem_next(context, builder, nextcontenttpe, tailtpe, nextcontentval, tailval, nextadvanced) + + outtpe = awkward1._numba.listoffsetarray.ListOffsetArrayType(util.indextpe(arraytpe.bitwidth), outcontenttpe, arraytpe.idtpe) + proxyout = numba.cgutils.create_struct_proxy(outtpe)(context, builder) + proxyout.offsets = nextoffsets + proxyout.content = outcontentval + if arraytpe.idtpe != numba.none: + proxyout.id = proxyin.id + return proxyout._getvalue() elif isinstance(headtpe, numba.types.EllipsisType): raise NotImplementedError("ListOffsetArray.getitem_next(ellipsis)") diff --git a/awkward1/_numba/util.py b/awkward1/_numba/util.py index 40b10a225f..da9a854e31 100644 --- a/awkward1/_numba/util.py +++ b/awkward1/_numba/util.py @@ -172,7 +172,7 @@ def regularize_slice(arrays): else: code = "def impl(arrays):\n" - indexes = [] + indexes = [] # all entries have trailing commas to ensure output is a tuple for i, t in enumerate(arrays.types): if isinstance(t, (numba.types.ArrayCompatible, numba.types.List)) and isinstance(t.dtype, numba.types.Boolean): code += " x{} = numpy.nonzero(arrays[{}])\n".format(i, i) @@ -182,7 +182,7 @@ def regularize_slice(arrays): elif isinstance(t, (numba.types.ArrayCompatible, numba.types.List)): raise TypeError("arrays must have boolean or integer type") else: - indexes.append("arrays[{}]".format(i)) + indexes.append("arrays[{}],".format(i)) code += " return ({})".format(" ".join(indexes)) g = {"numpy": numpy} exec(code, g) diff --git a/tests/test_PR015_finish_up_getitem_in_numba.py b/tests/test_PR015_finish_up_getitem_in_numba.py index 8d9b2ed50d..a9dbcd9317 100644 --- a/tests/test_PR015_finish_up_getitem_in_numba.py +++ b/tests/test_PR015_finish_up_getitem_in_numba.py @@ -139,6 +139,12 @@ def f3(q): assert awkward1.tolist(f3(listarrayB32)) == awkward1.tolist(f3(modelB)) + @numba.njit + def f4(q): + return q[numpy.array([0, 1, 1, 0]), ::2, 1:-1, 5::-1] + + assert awkward1.tolist(f4(listarrayB32)) == awkward1.tolist(f4(modelB)) + def test_deep_listoffsetarray32(): content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) @@ -168,6 +174,12 @@ def f3(q): assert awkward1.tolist(f3(listoffsetarrayB32)) == awkward1.tolist(f3(modelB)) + @numba.njit + def f4(q): + return q[numpy.array([0, 1, 1, 0]), ::2, 1:-1, 5::-1] + + assert awkward1.tolist(f4(listoffsetarrayB32)) == awkward1.tolist(f4(modelB)) + def test_deep_listarray64(): content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) @@ -199,6 +211,12 @@ def f3(q): assert awkward1.tolist(f3(listarrayB64)) == awkward1.tolist(f3(modelB)) + @numba.njit + def f4(q): + return q[numpy.array([0, 1, 1, 0]), ::2, 1:-1, 5::-1] + + assert awkward1.tolist(f4(listarrayB64)) == awkward1.tolist(f4(modelB)) + def test_deep_listoffsetarray64(): content = awkward1.layout.NumpyArray(numpy.arange(2*3*5*7).reshape(-1, 7)) offsetsA = numpy.arange(0, 2*3*5 + 5, 5) @@ -227,3 +245,9 @@ def f3(q): return q[1:, ::2, 1:-1, 5::-1] assert awkward1.tolist(f3(listoffsetarrayB64)) == awkward1.tolist(f3(modelB)) + + @numba.njit + def f4(q): + return q[numpy.array([0, 1, 1, 0]), ::2, 1:-1, 5::-1] + + assert awkward1.tolist(f4(listoffsetarrayB64)) == awkward1.tolist(f4(modelB)) From 70bd92659fe0d6362e1df5d4fb1ffaab60e67b3f Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Thu, 10 Oct 2019 09:50:59 -0500 Subject: [PATCH 36/36] Updated README; ready to squash and merge. --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index ed069bd28e..8629402df9 100644 --- a/README.md +++ b/README.md @@ -58,8 +58,9 @@ The following features of awkward 0.x will be features of awkward 1.x. * 2019-09-21 (PR [#9](../../pull/9)): `Identity` is correctly passed through `NumpyArray` slices and `__getitem__` uses `get`, `slice`, or the full `getitem`, depending on argument complexity. * 2019-09-26 (PR [#11](../../pull/11)): fully implemented `ListArray` and `ListOffsetArray`'s `__getitem__`. * 2019-10-02 (PR [#12](../../pull/12)): implemented `ListArray.__getitem__(array)` in Numba, setting the pattern for all the other cases. - * 2019-10-04 (PR [#13](../../pull13)): turned the `Error` type into a struct and `handle_error` into a function that provides `Identity` (if available) and attempted index information. Numba also handles errors, but with messages that must be known at compile-time. - * 2019-10-04 (PR [#14](../../pull14)): brought `NumpyArray`, `ListArray`, and `ListOffsetArray` up-to-date on `getitem` for all cases (except `newaxis`) and tested them all as well. (Not including Numba.) + * 2019-10-04 (PR [#13](../../pull/13)): turned the `Error` type into a struct and `handle_error` into a function that provides `Identity` (if available) and attempted index information. Numba also handles errors, but with messages that must be known at compile-time. + * 2019-10-04 (PR [#14](../../pull/14)): brought `NumpyArray`, `ListArray`, and `ListOffsetArray` up-to-date on `getitem` for all cases (except `newaxis`) and tested them all as well. (Not including Numba.) + * 2019-10-10 (PR [#15](../../pull/15)): brought `NumpyArray`, `ListArray`, and `ListOffsetArray` up-to-date on `getitem` in Numba (except `newaxis` and `Ellipsis`). ## Roadmap