Skip to content

Commit

Permalink
feat: add awkward_NumpyArray_subrange_equal and awkward_NumpyArray_su…
Browse files Browse the repository at this point in the history 
…brange_equal_bool kernel
  • Loading branch information
@ManasviGoyal
ManasviGoyal committed Jun 5, 2024
1 parent 0b9f6f4 commit 1ab7c7e
Show file tree
Hide file tree
Showing 6 changed files with 147 additions and 10 deletions.
2 changes: 2 additions & 0 deletions dev/generate-kernel-signatures.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,8 @@
"awkward_Content_getitem_next_missing_jagged_getmaskstartstop",
"awkward_index_rpad_and_clip_axis0",
"awkward_index_rpad_and_clip_axis1",
"awkward_NumpyArray_subrange_equal",
"awkward_NumpyArray_subrange_equal_bool",
"awkward_IndexedArray_flatten_nextcarry",
"awkward_IndexedArray_flatten_none2empty",
"awkward_IndexedArray_getitem_nextcarry",
Expand Down
2 changes: 2 additions & 0 deletions dev/generate-tests.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -867,6 +867,8 @@ def gencpuunittests(specdict):
"awkward_Content_getitem_next_missing_jagged_getmaskstartstop",
"awkward_index_rpad_and_clip_axis0",
"awkward_index_rpad_and_clip_axis1",
"awkward_NumpyArray_subrange_equal",
"awkward_NumpyArray_subrange_equal_bool",
"awkward_IndexedArray_flatten_nextcarry",
"awkward_IndexedArray_flatten_none2empty",
"awkward_IndexedArray_getitem_nextcarry",
Expand Down
4 changes: 2 additions & 2 deletions kernel-specification.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -543,7 +543,7 @@ kernels:
k = k + 1
automatic-tests: true

- name: awkward_IndexedArray_local_preparenext
- name: awkward_IndexedArray_local_preparenext_64
specializations:
- name: awkward_IndexedArray_local_preparenext_64
args:
Expand All @@ -555,7 +555,7 @@ kernels:
- {name: nextlen, type: "Const[int64_t]", dir: in, role: default}
description: null
definition: |
def awkward_IndexedArray_local_preparenext(
def awkward_IndexedArray_local_preparenext_64(
tocarry, starts, parents, parentslength, nextparents, nextlen
):
j = 0
Expand Down
138 changes: 134 additions & 4 deletions kernel-test-data.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16614,7 +16614,7 @@
]
},
{
"name": "awkward_IndexedArray_local_preparenext",
"name": "awkward_IndexedArray_local_preparenext_64",
"status": true,
"tests": [
{
Expand All @@ -16639,7 +16639,7 @@
"nextparents": [0, 0, 0, 0, 1, 1, 1],
"parentslength": 11,
"parents": [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
"starts": [0, 2, 5]
"starts": [0, 6]
},
"outputs": {
"tocarry": [0, 1, 2, 3, -1, -1, 4, 5, 6, -1, -1]
Expand Down Expand Up @@ -26387,12 +26387,38 @@
"tmpptr": [0, 2, 2, 3, 5],
"fromstarts": [0, 2, 3, 3],
"fromstops": [2, 3, 3, 5],
"length": 4
"length": 5
},
"outputs": {
"toequal": [0]
}
},
{
"error": false,
"message": "",
"inputs": {
"tmpptr": [0, 2, 2, 0, 2],
"fromstarts": [0, 2, 3, 3],
"fromstops": [2, 3, 3, 5],
"length": 5
},
"outputs": {
"toequal": [1]
}
},
{
"error": false,
"message": "",
"inputs": {
"tmpptr": [0, 0, 0, 0, 0],
"fromstarts": [0, 2, 3, 3],
"fromstops": [2, 3, 3, 5],
"length": 5
},
"outputs": {
"toequal": [1]
}
},
{
"error": false,
"message": "",
Expand All @@ -26406,6 +26432,32 @@
"toequal": [1]
}
},
{
"error": false,
"message": "",
"inputs": {
"tmpptr": [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1],
"fromstarts": [0, 2, 4, 6, 8, 10],
"fromstops": [2, 4, 6, 8, 10, 12],
"length": 6
},
"outputs": {
"toequal": [1]
}
},
{
"error": false,
"message": "",
"inputs": {
"tmpptr": [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 2],
"fromstarts": [0, 2, 4, 6, 8, 10],
"fromstops": [2, 4, 6, 8, 10, 12],
"length": 6
},
"outputs": {
"toequal": [1]
}
},
{
"error": false,
"message": "",
Expand All @@ -26418,6 +26470,19 @@
"outputs": {
"toequal": [1]
}
},
{
"error": false,
"message": "",
"inputs": {
"tmpptr": [1, 2, 3, 4, 5, 6],
"fromstarts": [2, 2, 2, 2, 2, 2],
"fromstops": [4, 4, 4, 4, 4, 4],
"length": 6
},
"outputs": {
"toequal": [1]
}
}
]
},
Expand Down Expand Up @@ -26458,12 +26523,38 @@
"tmpptr": [0, 2, 2, 3, 5],
"fromstarts": [0, 2, 3, 3],
"fromstops": [2, 3, 3, 5],
"length": 4
"length": 5
},
"outputs": {
"toequal": [0]
}
},
{
"error": false,
"message": "",
"inputs": {
"tmpptr": [0, 2, 2, 0, 2],
"fromstarts": [0, 2, 3, 3],
"fromstops": [2, 3, 3, 5],
"length": 5
},
"outputs": {
"toequal": [1]
}
},
{
"error": false,
"message": "",
"inputs": {
"tmpptr": [0, 0, 0, 0, 0],
"fromstarts": [0, 2, 3, 3],
"fromstops": [2, 3, 3, 5],
"length": 5
},
"outputs": {
"toequal": [1]
}
},
{
"error": false,
"message": "",
Expand All @@ -26477,6 +26568,32 @@
"toequal": [1]
}
},
{
"error": false,
"message": "",
"inputs": {
"tmpptr": [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1],
"fromstarts": [0, 2, 4, 6, 8, 10],
"fromstops": [2, 4, 6, 8, 10, 12],
"length": 6
},
"outputs": {
"toequal": [1]
}
},
{
"error": false,
"message": "",
"inputs": {
"tmpptr": [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 2],
"fromstarts": [0, 2, 4, 6, 8, 10],
"fromstops": [2, 4, 6, 8, 10, 12],
"length": 6
},
"outputs": {
"toequal": [1]
}
},
{
"error": false,
"message": "",
Expand All @@ -26489,6 +26606,19 @@
"outputs": {
"toequal": [1]
}
},
{
"error": false,
"message": "",
"inputs": {
"tmpptr": [1, 2, 3, 4, 5, 6],
"fromstarts": [2, 2, 2, 2, 2, 2],
"fromstops": [4, 4, 4, 4, 4, 4],
"length": 6
},
"outputs": {
"toequal": [1]
}
}
]
},
Expand Down
3 changes: 2 additions & 1 deletion src/awkward/contents/indexedarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -838,9 +838,10 @@ def _unique(self, negaxis, starts, parents, outlength):
and parents.nplike is self._backend.index_nplike
and nextparents.nplike is self._backend.index_nplike
)

self._backend.maybe_kernel_error(
self._backend[
"awkward_IndexedArray_local_preparenext",
"awkward_IndexedArray_local_preparenext_64",
nextoutindex.dtype.type,
starts.dtype.type,
parents.dtype.type,
Expand Down
8 changes: 5 additions & 3 deletions src/awkward/contents/indexedoptionarray.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -983,9 +983,10 @@ def _unique(self, negaxis, starts, parents, outlength):
and parents.nplike is self._backend.index_nplike
and nextparents.nplike is self._backend.index_nplike
)

self._backend.maybe_kernel_error(
self._backend[
"awkward_IndexedArray_local_preparenext",
"awkward_IndexedArray_local_preparenext_64",
nextoutindex.dtype.type,
starts.dtype.type,
parents.dtype.type,
Expand Down Expand Up @@ -1258,9 +1259,10 @@ def _argsort_next(
and parents.nplike is self._backend.index_nplike
and nextparents.nplike is self._backend.index_nplike
)

self._backend.maybe_kernel_error(
self._backend[
"awkward_IndexedArray_local_preparenext",
"awkward_IndexedArray_local_preparenext_64",
nextoutindex.dtype.type,
starts.dtype.type,
parents.dtype.type,
Expand Down Expand Up @@ -1342,7 +1344,7 @@ def _sort_next(self, negaxis, starts, parents, outlength, ascending, stable):

self._backend.maybe_kernel_error(
self._backend[
"awkward_IndexedArray_local_preparenext",
"awkward_IndexedArray_local_preparenext_64",
nextoutindex.dtype.type,
starts.dtype.type,
parents.dtype.type,
Expand Down

0 comments on commit 1ab7c7e

Please sign in to comment.