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

[jpivarski]

To narrow in on #3213, I made this branch that has print statements in the relevant methods (the _getitem chain for purely ragged numerical arrays for the slice [:, 0]). I first verified that RecordArrays have nothing to do with it.

The code to run here is

import awkward as ak

o = ak.from_parquet("s.parquet").a
o2 = ak.to_backend(o, "cuda")

print("========================= o[:, 0]")
tmp = o[:, 0]

print("============================ o2[:, 0]")
tmp = o2[:, 0]

using s.parquet.zip from #3213 (comment).

Starting at the Content._getitem entry point (and disregarding some odd stuff that the CUDA version does before this point, which I don't understand), the output is:

========================= o[:, 0]                                                                        ============================ o2[:, 0]
                                                                                                         
Content::_getitem <ListOffsetArray len='2000000'>                                                        Content::_getitem <ListOffsetArray len='2000000'>
    <offsets><Index dtype='int32' len='2000001'>                                                             <offsets><Index dtype='int32' len='2000001'>
        [      0       3       6 ... 3999998 3999999 4000000]                                                    [      0       3       6 ... 3999998 3999999 4000000]
    </Index></offsets>                                                                                       </Index></offsets>
    <content><ListOffsetArray len='4000000'>                                                                 <content><ListOffsetArray len='4000000'>
        <offsets><Index dtype='int32' len='4000001'>                                                             <offsets><Index dtype='int32' len='4000001'>
            [      0       3       3 ... 6999996 6999998 7000000]                                                    [      0       3       3 ... 6999996 6999998 7000000]
        </Index></offsets>                                                                                       </Index></offsets>
        <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>                <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>
    </ListOffsetArray></content>                                                                             </ListOffsetArray></content>
</ListOffsetArray> (slice(None, None, None), 0)                                                          </ListOffsetArray> (slice(None, None, None), 0)
                                                                                                         
ListOffsetArray::_carry <ListOffsetArray len='2000000'>                                                  ListOffsetArray::_carry <ListOffsetArray len='2000000'>
    <offsets><Index dtype='int32' len='2000001'>                                                             <offsets><Index dtype='int32' len='2000001'>
        [      0       3       6 ... 3999998 3999999 4000000]                                                    [      0       3       6 ... 3999998 3999999 4000000]
    </Index></offsets>                                                                                       </Index></offsets>
    <content><ListOffsetArray len='4000000'>                                                                 <content><ListOffsetArray len='4000000'>
        <offsets><Index dtype='int32' len='4000001'>                                                             <offsets><Index dtype='int32' len='4000001'>
            [      0       3       3 ... 6999996 6999998 7000000]                                                    [      0       3       3 ... 6999996 6999998 7000000]
        </Index></offsets>                                                                                       </Index></offsets>
        <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>                <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>
    </ListOffsetArray></content>                                                                             </ListOffsetArray></content>
</ListOffsetArray> <Index dtype='int64' len='2000000'>                                                   </ListOffsetArray> <Index dtype='int64' len='2000000'>
    [      0       1       2 ... 1999997 1999998 1999999]                                                    [      0       1       2 ... 1999997 1999998 1999999]
</Index> True                                                                                            </Index> True
                                                                                                         
ListArray::_getitem_next <ListArray len='2000000'>                                                       ListArray::_getitem_next <ListArray len='2000000'>
    <starts><Index dtype='int32' len='2000000'>                                                              <starts><Index dtype='int32' len='2000000'>
        [      0       3       6 ... 3999997 3999998 3999999]                                                    [      0       3       6 ... 3999997 3999998 3999999]
    </Index></starts>                                                                                        </Index></starts>
    <stops><Index dtype='int32' len='2000000'>                                                               <stops><Index dtype='int32' len='2000000'>
        [      3       6       9 ... 3999998 3999999 4000000]                                                    [      3       6       9 ... 3999998 3999999 4000000]
    </Index></stops>                                                                                         </Index></stops>
    <content><ListOffsetArray len='4000000'>                                                                 <content><ListOffsetArray len='4000000'>
        <offsets><Index dtype='int32' len='4000001'>                                                             <offsets><Index dtype='int32' len='4000001'>
            [      0       3       3 ... 6999996 6999998 7000000]                                                    [      0       3       3 ... 6999996 6999998 7000000]
        </Index></offsets>                                                                                       </Index></offsets>
        <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>                <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>
    </ListOffsetArray></content>                                                                             </ListOffsetArray></content>
</ListArray> 0 () None                                                                                   </ListArray> 0 () None
    lenstarts = 2000000                                                                                      lenstarts = 2000000
    self._starts = <Index dtype='int32' len='2000000'>                                                       self._starts = <Index dtype='int32' len='2000000'>
    [      0       3       6 ... 3999997 3999998 3999999]                                                    [      0       3       6 ... 3999997 3999998 3999999]
</Index>                                                                                                 </Index>
    self._stops = <Index dtype='int32' len='2000000'>                                                        self._stops = <Index dtype='int32' len='2000000'>
    [      3       6       9 ... 3999998 3999999 4000000]                                                    [      3       6       9 ... 3999998 3999999 4000000]
</Index>                                                                                                 </Index>
    nextcarry = <Index dtype='int64' len='2000000'>                                                          nextcarry = <Index dtype='int64' len='2000000'>
    [      0       3       6 ... 3999997 3999998 3999999]                                                    [125596686352384 125596686352387 125596686352390 ...               0
</Index>                                                                                                                    0               0]
                                                                                                         </Index>
ListOffsetArray::_carry <ListOffsetArray len='4000000'>                                                  
    <offsets><Index dtype='int32' len='4000001'>                                                         ListOffsetArray::_carry <ListOffsetArray len='4000000'>
        [      0       3       3 ... 6999996 6999998 7000000]                                                <offsets><Index dtype='int32' len='4000001'>
    </Index></offsets>                                                                                           [      0       3       3 ... 6999996 6999998 7000000]
    <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>                </Index></offsets>
</ListOffsetArray> <Index dtype='int64' len='2000000'>                                                       <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>
    [      0       3       6 ... 3999997 3999998 3999999]                                                </ListOffsetArray> <Index dtype='int64' len='2000000'>
</Index> True                                                                                                [125596686352384 125596686352387 125596686352390 ...               0
                                                                                                                            0               0]
ListArray::_getitem_next <ListArray len='2000000'>                                                       </Index> True
    <starts><Index dtype='int32' len='2000000'>                                                          
        [      0       5      10 ... 6999994 6999996 6999998]                                            ListArray::_getitem_next <ListArray len='2000000'>
    </Index></starts>                                                                                        <starts><Index dtype='int32' len='2000000'>
    <stops><Index dtype='int32' len='2000000'>                                                                   [3920640 3920645 3920650 ...       0       0       0]
        [      3       8      13 ... 6999996 6999998 7000000]                                                </Index></starts>
    </Index></stops>                                                                                         <stops><Index dtype='int32' len='2000000'>
    <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>                    [3920643 3920648 3920653 ...       3       3       3]
</ListArray> awkward._slicing.NO_HEAD () None                                                                </Index></stops>
                                                                                                             <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>
ListArray::_getitem_range <ListArray len='2000000'>                                                      </ListArray> awkward._slicing.NO_HEAD () None
    <starts><Index dtype='int32' len='2000000'>                                                          
        [      0       5      10 ... 6999994 6999996 6999998]                                            ListArray::_getitem_range <ListArray len='2000000'>
    </Index></starts>                                                                                        <starts><Index dtype='int32' len='2000000'>
    <stops><Index dtype='int32' len='2000000'>                                                                   [3920640 3920645 3920650 ...       0       0       0]
        [      3       8      13 ... 6999996 6999998 7000000]                                                </Index></starts>
    </Index></stops>                                                                                         <stops><Index dtype='int32' len='2000000'>
    <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>                    [3920643 3920648 3920653 ...       3       3       3]
</ListArray> 0 2000000                                                                                       </Index></stops>
                                                                                                             <content><NumpyArray dtype='int64' len='7000000'>[1 2 3 ... 7 6 7]</NumpyArray></content>
                                                                                                         </ListArray> 0 2000000

They diverge here:

awkward/src/awkward/contents/listarray.py

Lines 731 to 754 in 3ad5e20

	print(f" {lenstarts = }") # noqa: T201
	print(f" {self._starts = }") # noqa: T201
	print(f" {self._stops = }") # noqa: T201
	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,
	)
	print(f" {nextcarry = }") # noqa: T201

The CPU implementation of awkward_ListArray_getitem_next_at is okay, but what comes out of the CUDA implementation is random junk. Since CuPy wraps bad indexes in its implementation of fancy indexing, we don't get a segfault, just wrong answers.

Looking at awkward_ListArray_getitem_next_at.cpp and awkward_ListArray_getitem_next_at.cu, I don't immediately see the error—this ought to be a simple algorithm to parallelize. Maybe the issue is related to how it's called somehow? Is it missing a syncthreads to ensure that the result is assigned before moving on to the next step?

[ianna]

@jpivarski - I think, you can close the PR as the fix for the issue #3213 is provided by #3248