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!

dev/generate-kernel-signatures.py

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.

jpivarski

We talked about this at length in our meeting, and I'm on board with making this kind of change.

We also said that we'll need an automated check across all CUDA kernels (implemented once, where the kernels are built) to discover these C type errors in the future: int vs int*. (This is a missing feature of cupy.RawKernel!)

martindurant · 2024-09-27T20:49:44Z

I am still getting a breakage with selecting by integer, as opposed to slice (and hereafter, results on the GPU are garbage). I haven't managed to make a reproducer yet!

ianna · 2024-09-30T07:13:22Z

I am still getting a breakage with selecting by integer, as opposed to slice (and hereafter, results on the GPU are garbage). I haven't managed to make a reproducer yet!

hmm... I think, that what we have discussed with @jpivarski might happen.

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

Closed

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

dev/generate-kernel-signatures.py Show resolved Hide resolved

src/awkward/contents/regulararray.py Show resolved Hide resolved

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

ianna mentioned this pull request Sep 26, 2024

DO NOT MERGE: print-outs to diagnose issue #3213 #3216

Closed

ianna self-assigned this Sep 26, 2024

jpivarski approved these changes Sep 26, 2024

View reviewed changes

ianna merged commit e946646 into main Sep 26, 2024
44 checks passed

ianna deleted the 3213-bad-cuda-kernel-during-slice branch September 26, 2024 16:01

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

jpivarski commented Sep 20, 2024

ianna commented Sep 23, 2024

jpivarski left a comment

martindurant commented Sep 27, 2024

ianna commented Sep 30, 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

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

jpivarski commented Sep 20, 2024

ianna commented Sep 23, 2024

jpivarski left a comment

Choose a reason for hiding this comment

martindurant commented Sep 27, 2024

ianna commented Sep 30, 2024