fix: ListArray slicing on GPU #3248

ianna · 2024-09-18T18:37:24Z

No description provided.

ianna

@jpivarski - found the issue :-)
The numpy.ndarray head that is 0 is not passed to the GPU kernel!

awkward/src/awkward/contents/listarray.py

Line 731 in ba9b4d1

head,

in

awkward/src/awkward/contents/listarray.py

Lines 710 to 733 in ba9b4d1

    
           elif is_integer_like(head): 
        
               assert advanced is None 
        
               nexthead, nexttail = ak._slicing.head_tail(tail) 
        
               lenstarts = self._starts.length 
        
               nextcarry = ak.index.Index64.empty(lenstarts, self._backend.index_nplike) 
        
               assert ( 
        
                   nextcarry.nplike is self._backend.index_nplike 
        
                   and self._starts.nplike is self._backend.index_nplike 
        
                   and self._stops.nplike is self._backend.index_nplike 
        
               ) 
        
               self._maybe_index_error( 
        
                   self._backend[ 
        
                       "awkward_ListArray_getitem_next_at", 
        
                       nextcarry.dtype.type, 
        
                       self._starts.dtype.type, 
        
                       self._stops.dtype.type, 
        
                   ]( 
        
                       nextcarry.data, 
        
                       self._starts.data, 
        
                       self._stops.data, 
        
                       lenstarts, 
        
                       head, 
        
                   ), 
        
                   slicer=head,

head in this case can be an array and it can be regularized to a proper backend, then the GPU kernel needs to be updated to handle a 'cp.array(0)'

ianna

@martindurant - as discussed at today's awkward-uproot meeting this is a temporary fix. I will rewrite the way we handle this CUDA kernel wrapping it in a Python function and handling the 'at' type correctly.

martindurant · 2024-09-19T16:29:25Z

+1, please let me know when this is out.

martindurant · 2024-09-19T16:46:18Z

Any chance of a few more test cases like [:, 1:], [:, :1], [:, 1::2], [:, ::-1] ?

ianna

@jpivarski - Please, check when you have time. I think, using roles in this case is appropriate. I cannot find tests where the awkward_ListArray_getitem_jagged_expand is used yet. It looks like it should be invoked when a ListArray sliced with ListOffsetArray. It could be a different PR though. Thanks!

ianna · 2024-09-20T15:29:47Z

dev/generate-kernel-signatures.py

@@ -429,7 +429,10 @@ def by_signature(cuda_kernel_templates):
                special = [repr(spec["name"])]
                [type_to_pytype(x["type"], special) for x in childfunc["args"]]
                dirlist = [repr(x["dir"]) for x in childfunc["args"]]
-                ispointerlist = [repr("List" in x["type"]) for x in childfunc["args"]]
+                ispointerlist = [
+                    repr("List" in x["type"] or "ListArray-at" == x.get("role", None))


This is a first use of 'role' to specify that the type is a pointer in a cuda kernel. There are three kernels that use this role: two of them are in this PR and the third one is awkward_ListArray_getitem_jagged_expand.

Right now, the roles are like untested documentation: I don't think anything relies on them (except this ispointerlist, now), so they might be inaccurate or not correspond to the ispointerlist condition that you want here.

The ListArray-at role is not defined nor used anywhere else. This is it's first usage definition. Let's discuss it on Wednesday.

src/awkward/contents/regulararray.py

jpivarski · 2024-09-20T15:49:42Z

I cannot find tests where the awkward_ListArray_getitem_jagged_expand is used yet.

I used the trick of adding

        if index[0] == "awkward_ListArray_getitem_jagged_expand":
            raise Exception("HERE")

to

awkward/src/awkward/_backends/numpy.py

Lines 35 to 36 in 3712032

    
           def __getitem__(self, index: KernelKeyType) -> NumpyKernel: 
        
               return NumpyKernel(awkward_cpp.cpu_kernels.kernel[index], index)

and found exactly 1 test where it's used: tests/test_0111_jagged_and_masked_getitem.py in test_double_jagged.

    def test_double_jagged():
        array = ak.highlevel.Array(
            [[[0, 1, 2, 3], [4, 5]], [[6, 7, 8], [9, 10, 11, 12, 13]]], check_valid=True
        ).layout
        array2 = ak.highlevel.Array(
            [[[2, 1, 0], [-1]], [[-1, -2, -3], [2, 1, 1, 3]]], check_valid=True
        ).layout
    
        assert to_list(array[array2]) == [
            [[2, 1, 0], [5]],
            [[8, 7, 6], [11, 10, 10, 12]],
        ]
        assert array.to_typetracer()[array2].form == array[array2].form
    
        content = ak.operations.from_iter(
            [[0, 1, 2, 3], [4, 5], [6, 7, 8], [9, 10, 11, 12, 13]], highlevel=False
        )
        regulararray = ak.contents.RegularArray(content, 2, zeros_length=0)
    
        array1 = ak.highlevel.Array([[2, 1, 0], [-1]], check_valid=True).layout
    
        assert to_list(regulararray[:, array1]) == [[[2, 1, 0], [5]], [[8, 7, 6], [13]]]
        assert regulararray.to_typetracer()[:, array1].form == regulararray[:, array1].form
        assert to_list(regulararray[1:, array1]) == [[[8, 7, 6], [13]]]
        assert (
            regulararray.to_typetracer()[1:, array1].form == regulararray[1:, array1].form
        )
    
        offsets = ak.index.Index64(np.array([0, 2, 4], dtype=np.int64))
        listoffsetarray = ak.contents.ListOffsetArray(offsets, content)
>       assert to_list(listoffsetarray[:, array1]) == [
            [[2, 1, 0], [5]],
            [[8, 7, 6], [13]],
        ]

array      = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
array1     = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>
array2     = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte... [ 2  1  0 -1 -1 -2 -3  2  1  1  3]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
content    = <ListOffsetArray len='4'>
    <offsets><Index dtype='int64' len='5'>
        [ 0  4  6  9 14]
    </Index></offsets>
 ...pe='int64' len='14'>
        [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13]
    </NumpyArray></content>
</ListOffsetArray>
listoffsetarray = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
offsets    = <Index dtype='int64' len='3'>
    [0 2 4]
</Index>
regulararray = <RegularArray size='2' len='2'>
    <content><ListOffsetArray len='4'>
        <offsets><Index dtype='int64' len='5'>
...1  2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</RegularArray>

tests/test_0111_jagged_and_masked_getitem.py:711: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
src/awkward/contents/content.py:512: in __getitem__
    return self._getitem(where)
        self       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
        where      = (slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>)
src/awkward/contents/content.py:557: in _getitem
    out = next._getitem_next(nextwhere[0], nextwhere[1:], None)
        backend    = <awkward._backends.numpy.NumpyBackend object at 0x783910736210>
        items      = [slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>]
        next       = <RegularArray size='2' len='1'>
    <content><ListOffsetArray len='2'>
        <offsets><Index dtype='int64' len='3'>
...          </NumpyArray></content>
        </ListOffsetArray></content>
    </ListOffsetArray></content>
</RegularArray>
        nextwhere  = (slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>)
        self       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
        this       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
        where      = (slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>)
src/awkward/contents/regulararray.py:519: in _getitem_next
    nextcontent._getitem_next(nexthead, nexttail, advanced),
        advanced   = None
        head       = slice(None, None, None)
        index_nplike = <awkward._nplikes.numpy.Numpy object at 0x783910703f50>
        nextcarry  = <Index dtype='int64' len='2'>[0 1]</Index>
        nextcontent = <ListArray len='2'>
    <starts><Index dtype='int64' len='2'>
        [0 2]
    </Index></starts>
    <stops><Index dt...0  1  2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListArray>
        nexthead   = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>
        nextsize   = 2
        nexttail   = ()
        self       = <RegularArray size='2' len='1'>
    <content><ListOffsetArray len='2'>
        <offsets><Index dtype='int64' len='3'>
...          </NumpyArray></content>
        </ListOffsetArray></content>
    </ListOffsetArray></content>
</RegularArray>
        start      = 0
        step       = 1
        stop       = 2
        tail       = (<ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>,)
src/awkward/contents/listarray.py:1026: in _getitem_next
    self._backend[
        advanced   = None
        head       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>
        headlength = 2
        length     = 2
        multistarts = <Index dtype='int64' len='4'>[ 8  7  6 13]</Index>
        multistops = <Index dtype='int64' len='4'>[ 8  7  6 13]</Index>
        nextcarry  = <Index dtype='int64' len='4'>[3 4 3 4]</Index>
        self       = <ListArray len='2'>
    <starts><Index dtype='int64' len='2'>
        [0 2]
    </Index></starts>
    <stops><Index dt...0  1  2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListArray>
        singleoffsets = <Index dtype='int64' len='3'>[0 3 4]</Index>
        tail       = ()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <awkward._backends.numpy.NumpyBackend object at 0x783910736210>
index = ('awkward_ListArray_getitem_jagged_expand', <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, ...)

    def __getitem__(self, index: KernelKeyType) -> NumpyKernel:
        if index[0] == "awkward_ListArray_getitem_jagged_expand":
>           raise Exception("HERE")
E           Exception: HERE

index      = ('awkward_ListArray_getitem_jagged_expand', <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, ...)
self       = <awkward._backends.numpy.NumpyBackend object at 0x783910736210>

src/awkward/_backends/numpy.py:37: Exception

ianna · 2024-09-23T13:28:32Z

I cannot find tests where the awkward_ListArray_getitem_jagged_expand is used yet.

I used the trick of adding

        if index[0] == "awkward_ListArray_getitem_jagged_expand":
            raise Exception("HERE")

to

awkward/src/awkward/_backends/numpy.py

Lines 35 to 36 in 3712032

    
           def __getitem__(self, index: KernelKeyType) -> NumpyKernel: 
        
               return NumpyKernel(awkward_cpp.cpu_kernels.kernel[index], index)

and found exactly 1 test where it's used: tests/test_0111_jagged_and_masked_getitem.py in test_double_jagged.

    def test_double_jagged():
        array = ak.highlevel.Array(
            [[[0, 1, 2, 3], [4, 5]], [[6, 7, 8], [9, 10, 11, 12, 13]]], check_valid=True
        ).layout
        array2 = ak.highlevel.Array(
            [[[2, 1, 0], [-1]], [[-1, -2, -3], [2, 1, 1, 3]]], check_valid=True
        ).layout
    
        assert to_list(array[array2]) == [
            [[2, 1, 0], [5]],
            [[8, 7, 6], [11, 10, 10, 12]],
        ]
        assert array.to_typetracer()[array2].form == array[array2].form
    
        content = ak.operations.from_iter(
            [[0, 1, 2, 3], [4, 5], [6, 7, 8], [9, 10, 11, 12, 13]], highlevel=False
        )
        regulararray = ak.contents.RegularArray(content, 2, zeros_length=0)
    
        array1 = ak.highlevel.Array([[2, 1, 0], [-1]], check_valid=True).layout
    
        assert to_list(regulararray[:, array1]) == [[[2, 1, 0], [5]], [[8, 7, 6], [13]]]
        assert regulararray.to_typetracer()[:, array1].form == regulararray[:, array1].form
        assert to_list(regulararray[1:, array1]) == [[[8, 7, 6], [13]]]
        assert (
            regulararray.to_typetracer()[1:, array1].form == regulararray[1:, array1].form
        )
    
        offsets = ak.index.Index64(np.array([0, 2, 4], dtype=np.int64))
        listoffsetarray = ak.contents.ListOffsetArray(offsets, content)
>       assert to_list(listoffsetarray[:, array1]) == [
            [[2, 1, 0], [5]],
            [[8, 7, 6], [13]],
        ]

array      = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
array1     = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>
array2     = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte... [ 2  1  0 -1 -1 -2 -3  2  1  1  3]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
content    = <ListOffsetArray len='4'>
    <offsets><Index dtype='int64' len='5'>
        [ 0  4  6  9 14]
    </Index></offsets>
 ...pe='int64' len='14'>
        [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13]
    </NumpyArray></content>
</ListOffsetArray>
listoffsetarray = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
offsets    = <Index dtype='int64' len='3'>
    [0 2 4]
</Index>
regulararray = <RegularArray size='2' len='2'>
    <content><ListOffsetArray len='4'>
        <offsets><Index dtype='int64' len='5'>
...1  2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</RegularArray>

tests/test_0111_jagged_and_masked_getitem.py:711: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
src/awkward/contents/content.py:512: in __getitem__
    return self._getitem(where)
        self       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
        where      = (slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>)
src/awkward/contents/content.py:557: in _getitem
    out = next._getitem_next(nextwhere[0], nextwhere[1:], None)
        backend    = <awkward._backends.numpy.NumpyBackend object at 0x783910736210>
        items      = [slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>]
        next       = <RegularArray size='2' len='1'>
    <content><ListOffsetArray len='2'>
        <offsets><Index dtype='int64' len='3'>
...          </NumpyArray></content>
        </ListOffsetArray></content>
    </ListOffsetArray></content>
</RegularArray>
        nextwhere  = (slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>)
        self       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
        this       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
        where      = (slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>)
src/awkward/contents/regulararray.py:519: in _getitem_next
    nextcontent._getitem_next(nexthead, nexttail, advanced),
        advanced   = None
        head       = slice(None, None, None)
        index_nplike = <awkward._nplikes.numpy.Numpy object at 0x783910703f50>
        nextcarry  = <Index dtype='int64' len='2'>[0 1]</Index>
        nextcontent = <ListArray len='2'>
    <starts><Index dtype='int64' len='2'>
        [0 2]
    </Index></starts>
    <stops><Index dt...0  1  2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListArray>
        nexthead   = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>
        nextsize   = 2
        nexttail   = ()
        self       = <RegularArray size='2' len='1'>
    <content><ListOffsetArray len='2'>
        <offsets><Index dtype='int64' len='3'>
...          </NumpyArray></content>
        </ListOffsetArray></content>
    </ListOffsetArray></content>
</RegularArray>
        start      = 0
        step       = 1
        stop       = 2
        tail       = (<ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>,)
src/awkward/contents/listarray.py:1026: in _getitem_next
    self._backend[
        advanced   = None
        head       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>
        headlength = 2
        length     = 2
        multistarts = <Index dtype='int64' len='4'>[ 8  7  6 13]</Index>
        multistops = <Index dtype='int64' len='4'>[ 8  7  6 13]</Index>
        nextcarry  = <Index dtype='int64' len='4'>[3 4 3 4]</Index>
        self       = <ListArray len='2'>
    <starts><Index dtype='int64' len='2'>
        [0 2]
    </Index></starts>
    <stops><Index dt...0  1  2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListArray>
        singleoffsets = <Index dtype='int64' len='3'>[0 3 4]</Index>
        tail       = ()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <awkward._backends.numpy.NumpyBackend object at 0x783910736210>
index = ('awkward_ListArray_getitem_jagged_expand', <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, ...)

    def __getitem__(self, index: KernelKeyType) -> NumpyKernel:
        if index[0] == "awkward_ListArray_getitem_jagged_expand":
>           raise Exception("HERE")
E           Exception: HERE

index      = ('awkward_ListArray_getitem_jagged_expand', <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, ...)
self       = <awkward._backends.numpy.NumpyBackend object at 0x783910736210>

src/awkward/_backends/numpy.py:37: Exception

Thanks! Yes, this test is on

I cannot find tests where the awkward_ListArray_getitem_jagged_expand is used yet.

I used the trick of adding

        if index[0] == "awkward_ListArray_getitem_jagged_expand":
            raise Exception("HERE")

to

awkward/src/awkward/_backends/numpy.py

Lines 35 to 36 in 3712032

    
           def __getitem__(self, index: KernelKeyType) -> NumpyKernel: 
        
               return NumpyKernel(awkward_cpp.cpu_kernels.kernel[index], index)

and found exactly 1 test where it's used: tests/test_0111_jagged_and_masked_getitem.py in test_double_jagged.

    def test_double_jagged():
        array = ak.highlevel.Array(
            [[[0, 1, 2, 3], [4, 5]], [[6, 7, 8], [9, 10, 11, 12, 13]]], check_valid=True
        ).layout
        array2 = ak.highlevel.Array(
            [[[2, 1, 0], [-1]], [[-1, -2, -3], [2, 1, 1, 3]]], check_valid=True
        ).layout
    
        assert to_list(array[array2]) == [
            [[2, 1, 0], [5]],
            [[8, 7, 6], [11, 10, 10, 12]],
        ]
        assert array.to_typetracer()[array2].form == array[array2].form
    
        content = ak.operations.from_iter(
            [[0, 1, 2, 3], [4, 5], [6, 7, 8], [9, 10, 11, 12, 13]], highlevel=False
        )
        regulararray = ak.contents.RegularArray(content, 2, zeros_length=0)
    
        array1 = ak.highlevel.Array([[2, 1, 0], [-1]], check_valid=True).layout
    
        assert to_list(regulararray[:, array1]) == [[[2, 1, 0], [5]], [[8, 7, 6], [13]]]
        assert regulararray.to_typetracer()[:, array1].form == regulararray[:, array1].form
        assert to_list(regulararray[1:, array1]) == [[[8, 7, 6], [13]]]
        assert (
            regulararray.to_typetracer()[1:, array1].form == regulararray[1:, array1].form
        )
    
        offsets = ak.index.Index64(np.array([0, 2, 4], dtype=np.int64))
        listoffsetarray = ak.contents.ListOffsetArray(offsets, content)
>       assert to_list(listoffsetarray[:, array1]) == [
            [[2, 1, 0], [5]],
            [[8, 7, 6], [13]],
        ]

array      = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
array1     = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>
array2     = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte... [ 2  1  0 -1 -1 -2 -3  2  1  1  3]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
content    = <ListOffsetArray len='4'>
    <offsets><Index dtype='int64' len='5'>
        [ 0  4  6  9 14]
    </Index></offsets>
 ...pe='int64' len='14'>
        [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13]
    </NumpyArray></content>
</ListOffsetArray>
listoffsetarray = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
offsets    = <Index dtype='int64' len='3'>
    [0 2 4]
</Index>
regulararray = <RegularArray size='2' len='2'>
    <content><ListOffsetArray len='4'>
        <offsets><Index dtype='int64' len='5'>
...1  2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</RegularArray>

tests/test_0111_jagged_and_masked_getitem.py:711: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
src/awkward/contents/content.py:512: in __getitem__
    return self._getitem(where)
        self       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
        where      = (slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>)
src/awkward/contents/content.py:557: in _getitem
    out = next._getitem_next(nextwhere[0], nextwhere[1:], None)
        backend    = <awkward._backends.numpy.NumpyBackend object at 0x783910736210>
        items      = [slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>]
        next       = <RegularArray size='2' len='1'>
    <content><ListOffsetArray len='2'>
        <offsets><Index dtype='int64' len='3'>
...          </NumpyArray></content>
        </ListOffsetArray></content>
    </ListOffsetArray></content>
</RegularArray>
        nextwhere  = (slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>)
        self       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
        this       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 2 4]
    </Index></offsets>
    <conte...2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListOffsetArray>
        where      = (slice(None, None, None), <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>)
src/awkward/contents/regulararray.py:519: in _getitem_next
    nextcontent._getitem_next(nexthead, nexttail, advanced),
        advanced   = None
        head       = slice(None, None, None)
        index_nplike = <awkward._nplikes.numpy.Numpy object at 0x783910703f50>
        nextcarry  = <Index dtype='int64' len='2'>[0 1]</Index>
        nextcontent = <ListArray len='2'>
    <starts><Index dtype='int64' len='2'>
        [0 2]
    </Index></starts>
    <stops><Index dt...0  1  2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListArray>
        nexthead   = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>
        nextsize   = 2
        nexttail   = ()
        self       = <RegularArray size='2' len='1'>
    <content><ListOffsetArray len='2'>
        <offsets><Index dtype='int64' len='3'>
...          </NumpyArray></content>
        </ListOffsetArray></content>
    </ListOffsetArray></content>
</RegularArray>
        start      = 0
        step       = 1
        stop       = 2
        tail       = (<ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>,)
src/awkward/contents/listarray.py:1026: in _getitem_next
    self._backend[
        advanced   = None
        head       = <ListOffsetArray len='2'>
    <offsets><Index dtype='int64' len='3'>
        [0 3 4]
    </Index></offsets>
    <content><NumpyArray dtype='int64' len='4'>[ 2  1  0 -1]</NumpyArray></content>
</ListOffsetArray>
        headlength = 2
        length     = 2
        multistarts = <Index dtype='int64' len='4'>[ 8  7  6 13]</Index>
        multistops = <Index dtype='int64' len='4'>[ 8  7  6 13]</Index>
        nextcarry  = <Index dtype='int64' len='4'>[3 4 3 4]</Index>
        self       = <ListArray len='2'>
    <starts><Index dtype='int64' len='2'>
        [0 2]
    </Index></starts>
    <stops><Index dt...0  1  2  3  4  5  6  7  8  9 10 11 12 13]
        </NumpyArray></content>
    </ListOffsetArray></content>
</ListArray>
        singleoffsets = <Index dtype='int64' len='3'>[0 3 4]</Index>
        tail       = ()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <awkward._backends.numpy.NumpyBackend object at 0x783910736210>
index = ('awkward_ListArray_getitem_jagged_expand', <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, ...)

    def __getitem__(self, index: KernelKeyType) -> NumpyKernel:
        if index[0] == "awkward_ListArray_getitem_jagged_expand":
>           raise Exception("HERE")
E           Exception: HERE

index      = ('awkward_ListArray_getitem_jagged_expand', <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, <class 'numpy.int64'>, ...)
self       = <awkward._backends.numpy.NumpyBackend object at 0x783910736210>

src/awkward/_backends/numpy.py:37: Exception

The CUDA tests are implemented for this and the results are equivalent to the ones produced on CPU.

test: add slicing test for CPU and GPU in test_3140_cuda_slicing.py

1b826c9

ianna linked an issue Sep 18, 2024 that may be closed by this pull request

Bad CUDA kernel during slice #3213

Open

ianna marked this pull request as draft September 18, 2024 18:37

style: pre-commit fixes

ba9b4d1

ianna commented Sep 18, 2024

View reviewed changes

cast 'at' to int

ef2609e

head in this case can be an array and it can be regularized to a proper backend, then the GPU kernel needs to be updated to handle a 'cp.array(0)'

ianna marked this pull request as ready for review September 19, 2024 13:21

style: pre-commit fixes

ed55752

ianna requested a review from jpivarski September 19, 2024 13:22

pre-commit-ci bot temporarily deployed to docs September 19, 2024 13:33 Inactive

ianna marked this pull request as draft September 19, 2024 13:44

use ak._slicing.normalize_integer_like(head)

bd7b125

ianna temporarily deployed to docs September 19, 2024 14:04 — with GitHub Actions Inactive

convert head ndarray to scalar

298ecf9

ianna temporarily deployed to docs September 19, 2024 14:45 — with GitHub Actions Inactive

ianna added 3 commits September 19, 2024 16:52

add item attribute to TypeTracerArray

2627a51

cleanup tests

6333a6c

Merge branch 'main' into 3213-bad-cuda-kernel-during-slice

9bbd8f6

ianna marked this pull request as ready for review September 19, 2024 15:04

ianna temporarily deployed to docs September 19, 2024 15:07 — with GitHub Actions Inactive

ianna commented Sep 19, 2024

View reviewed changes

ianna and others added 2 commits September 19, 2024 19:03

add more tests

da434dc

style: pre-commit fixes

5ed2072

pre-commit-ci bot temporarily deployed to docs September 19, 2024 17:15 Inactive

ianna and others added 4 commits September 20, 2024 16:36

use 'ListArray-at' role

2fdbb2a

style: pre-commit fixes

095a2ed

use pointer awkward_ListArray_getitem_next_at.cu

408a24b

use role listarray.py

2a12275

revert changes

13d3a49

ianna marked this pull request as draft September 20, 2024 14:39

Merge branch 'main' into 3213-bad-cuda-kernel-during-slice

84d47a8

ianna temporarily deployed to docs September 20, 2024 14:52 — with GitHub Actions Inactive

ianna added 2 commits September 20, 2024 17:13

same for awkward_RegularArray_getitem_next_at.cu

11f3f4e

remove special case for cuda

a7f8f3b

ianna temporarily deployed to docs September 20, 2024 15:27 — with GitHub Actions Inactive

ianna commented Sep 20, 2024

View reviewed changes

ianna marked this pull request as ready for review September 20, 2024 15:35

ianna changed the title ~~fix: slicing on GPU~~ fix: ListArray slicing on GPU Sep 20, 2024

correct role for jagged size

a8845ee

ianna deployed to docs September 23, 2024 14:04 — with GitHub Actions View deployment

ianna marked this pull request as draft September 23, 2024 14:27

ianna marked this pull request as ready for review September 24, 2024 09:52

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: ListArray slicing on GPU #3248

fix: ListArray slicing on GPU #3248

ianna commented Sep 18, 2024

ianna left a comment

ianna left a comment

martindurant commented Sep 19, 2024

martindurant commented Sep 19, 2024

ianna left a comment

ianna Sep 20, 2024

jpivarski Sep 20, 2024

ianna Sep 23, 2024

jpivarski commented Sep 20, 2024

ianna commented Sep 23, 2024

	elif is_integer_like(head):
	assert advanced is None
	nexthead, nexttail = ak._slicing.head_tail(tail)
	lenstarts = self._starts.length
	nextcarry = ak.index.Index64.empty(lenstarts, self._backend.index_nplike)
	assert (
	nextcarry.nplike is self._backend.index_nplike
	and self._starts.nplike is self._backend.index_nplike
	and self._stops.nplike is self._backend.index_nplike
	)
	self._maybe_index_error(
	self._backend[
	"awkward_ListArray_getitem_next_at",
	nextcarry.dtype.type,
	self._starts.dtype.type,
	self._stops.dtype.type,
	](
	nextcarry.data,
	self._starts.data,
	self._stops.data,
	lenstarts,
	head,
	),
	slicer=head,

fix: ListArray slicing on GPU #3248

Are you sure you want to change the base?

fix: ListArray slicing on GPU #3248

Conversation

ianna commented Sep 18, 2024

ianna left a comment

Choose a reason for hiding this comment

ianna left a comment

Choose a reason for hiding this comment

martindurant commented Sep 19, 2024

martindurant commented Sep 19, 2024

ianna left a comment

Choose a reason for hiding this comment

ianna Sep 20, 2024

Choose a reason for hiding this comment

jpivarski Sep 20, 2024

Choose a reason for hiding this comment

ianna Sep 23, 2024

Choose a reason for hiding this comment

jpivarski commented Sep 20, 2024

ianna commented Sep 23, 2024