feat: add awkward_IndexedArray_numnull kernel implementation using cuda.compute

[maxymnaumchyk]

Closes #3957

[github-actions]

The documentation preview is ready to be viewed at http://preview.awkward-array.org.s3-website.us-east-1.amazonaws.com/PR3960

[codecov]

Codecov Report

❌ Patch coverage is 87.50000% with 1 line in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@80c3e39). Learn more about missing BASE report.
⚠️ Report is 41 commits behind head on main.

Files with missing lines	Patch %	Lines
src/awkward/_connect/cuda/_compute.py	85.71%	1 Missing ⚠️

❌ Your patch check has failed because the patch coverage (87.50%) is below the target coverage (98.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/_backends/cupy.py	100.00% <100.00%> (ø)
src/awkward/_connect/cuda/_compute.py	57.93% <85.71%> (ø)

[maxymnaumchyk]

Some profiling

Old implementation of awkward_IndexedArray_numnull

python test_IndexedArray_numnull.py

Time taken: 0.012456 seconds
Input data size: 524287996 bytes (500.0 MB)

Implementation of awkward_IndexedArray_numnull using cupy

python test_IndexedArray_numnull.py

Time taken: 0.001752 seconds
Input data size: 524287996 bytes (500.0 MB)

Also, as a note I had inconsistent measurements here which is probalbly because I'm using WSL (there was a discussion about that here)

test_IndexedArray_numnull.py Script code

import timeit

import cupy
import nvtx

from awkward._backends.cupy import CupyBackend

cupy_backend = CupyBackend.instance()

# ~500 MB: 7 elements * 18_724_571 repeats * 4 bytes = ~524 MB
REPEAT = 18_724_571
base = [0, -1, 2, -1, -1, -1, -1]

numnull = cupy.array([0], dtype=cupy.int64)
fromindex = cupy.array(base * REPEAT, dtype=cupy.int32)
lenindex = len(fromindex)
funcC = cupy_backend['awkward_IndexedArray_numnull', cupy.int64, cupy.int32]

# warmup
funcC(numnull, fromindex, lenindex)

# Benchmark
start_time = timeit.default_timer()
for _ in range(10):
    with nvtx.annotate("kernel call", color="blue"):
        funcC(numnull, fromindex, lenindex)
    cupy.cuda.Device().synchronize()
end_time = timeit.default_timer()

print(f"Time taken: {(end_time - start_time) / 10:.6f} seconds")

# Memory usage report
total_bytes = numnull.nbytes + fromindex.nbytes
print(f"Input data size: {total_bytes} bytes ({total_bytes / 1024**2:.1f} MB)")

[maxymnaumchyk]

Using cuda.compute

def awkward_IndexedArray_numnull(numnull, fromindex, lenindex):
    index_dtype = numnull.dtype

    def is_null(x):
        return 1 if x < 0 else 0

    null_iter = TransformIterator(fromindex[:lenindex], is_null)

    h_init = np.array([0], dtype=index_dtype)
    reduce_into(null_iter, numnull, OpKind.PLUS, lenindex, h_init)

The fastest run is ~2 times faster than using cupy. Inconsistency is probably because of the same reason (testing on WSL)

[maxymnaumchyk]

Using this script shows consistent performance measurements:

test_IndexedArray_numnull.py Script code

import timeit

import cupy
import nvtx

from awkward._backends.cupy import CupyBackend

cupy_backend = CupyBackend.instance()

REPEAT = 18000000
base = [0, -1, 2, -1, -1, -1, -1]

numnull = cupy.array([0], dtype=cupy.int64)
fromindex = cupy.array(base * REPEAT, dtype=cupy.int32)
lenindex = len(fromindex)
funcC = cupy_backend['awkward_IndexedArray_numnull', cupy.int64, cupy.int32]

# warmup
for _ in range(15):
    funcC(numnull, fromindex, lenindex)
cupy.cuda.Device().synchronize()

# Benchmark
n_iter = 50

start_time = timeit.default_timer()
with nvtx.annotate("benchmarking", color="red"):
    for _ in range(n_iter):
        with nvtx.annotate("kernel call", color="blue"):
            funcC(numnull, fromindex, lenindex)
    cupy.cuda.Device().synchronize()
end_time = timeit.default_timer()

print(f"Time taken: {(end_time - start_time) / n_iter:.6f} seconds")

# Memory usage report
total_bytes = numnull.nbytes + fromindex.nbytes
print(f"Input data size: {total_bytes} bytes ({total_bytes / 1024**2:.1f} MB)")

[ianna]

@maxymnaumchyk -- I actually do not see the awkward_IndexedArray_numnull implementation

[maxymnaumchyk]

I will close this PR as it was transferred here: #4124