fix: segmented_reduce should use max_segment_size

[maxymnaumchyk]

No description provided.

[maxymnaumchyk]

Using max_segment_size speeds things up as expected.

However if I try to compare the resulting kernel with the previous implementation (which was using unary_transform) it turns out to be slower. The reason for that is: although I am using max_segment_size I have to first calculate it with max_segment_size = cupy.max(end_o - start_o)

full script

import math
import timeit

import cupy
import numpy
import nvtx
from cuda.compute import CountingIterator, OpKind, segmented_reduce, unary_transform

# --- compile old CUDA kernel ---
NO_ERROR = numpy.iinfo(numpy.uint64).max
ERROR_BITS = 8

kernels_path = "src/awkward/_connect/cuda/cuda_kernels"
with open(f"{kernels_path}/cuda_common.cu") as f:
    common_src = f.read()
with open(f"{kernels_path}/awkward_reduce_prod.cu") as f:
    cu_src = f.read()

cuda_src = (
    f"#define ERROR_BITS {ERROR_BITS}\n#define NO_ERROR {NO_ERROR}ULL\n"
    + common_src
    + "\n"
    + cu_src
)

old_module = cupy.RawModule(
    code=cuda_src,
    options=("--std=c++11",),
    name_expressions=["awkward_reduce_prod_kernel<int64_t, int64_t, int64_t>"],
)
old_kernel = old_module.get_function(
    "awkward_reduce_prod_kernel<int64_t, int64_t, int64_t>"
)

# --- data setup ---
N_SEGMENTS = 1_000_000
SEGMENT_SIZE = 10

offsets = cupy.arange(
    0, (N_SEGMENTS + 1) * SEGMENT_SIZE, SEGMENT_SIZE, dtype=cupy.int64
)
input_data = cupy.ones(N_SEGMENTS * SEGMENT_SIZE, dtype=cupy.int64)
result = cupy.empty(N_SEGMENTS, dtype=cupy.int64)
start_o = offsets[:-1]
end_o = offsets[1:]
h_init = numpy.asarray(1, dtype=numpy.int64)
mss = int(cupy.max(end_o - start_o))

invocation_index = cupy.array(NO_ERROR, dtype=cupy.uint64)
err_code = cupy.array(NO_ERROR, dtype=cupy.uint64)

BLOCK_SIZE = 512


def run_old(result):
    grid_size = math.ceil(N_SEGMENTS / BLOCK_SIZE)
    old_kernel(
        (grid_size,),
        (BLOCK_SIZE,),
        (result, input_data, offsets, numpy.int64(N_SEGMENTS), invocation_index, err_code),
    )


def run_unary_transform(result):
    def segment_reduce_prod(segment_id):
        start_idx = start_o[segment_id]
        end_idx = end_o[segment_id]
        segment = input_data[start_idx:end_idx]
        if len(segment) == 0:
            return numpy.int64(1)
        return numpy.prod(segment)

    type_wrapper = cupy.dtype(cupy.int64).type
    segment_ids = CountingIterator(type_wrapper(0))
    unary_transform(d_in=segment_ids, d_out=result, op=segment_reduce_prod, num_items=N_SEGMENTS)


def run_new(result):
    segmented_reduce(
        d_in=input_data,
        d_out=result,
        num_segments=N_SEGMENTS,
        start_offsets_in=start_o,
        end_offsets_in=end_o,
        op=OpKind.PLUS,
        h_init=h_init,
        max_segment_size=mss,
    )


# warmup
for _ in range(15):
    run_old(result)
    run_unary_transform(result)
    run_new(result)
cupy.cuda.Device().synchronize()

n_iter = 50

# --- old kernel ---
start = timeit.default_timer()
with nvtx.annotate("old CUDA kernel", color="red"):
    for _ in range(n_iter):
        with nvtx.annotate("kernel call", color="blue"):
            run_old(result)
    cupy.cuda.Device().synchronize()
t_old = (timeit.default_timer() - start) / n_iter

# --- unary_transform ---
start = timeit.default_timer()
with nvtx.annotate("unary_transform", color="purple"):
    for _ in range(n_iter):
        with nvtx.annotate("kernel call", color="orange"):
            run_unary_transform(result)
    cupy.cuda.Device().synchronize()
t_unary = (timeit.default_timer() - start) / n_iter

# --- new cuda.compute kernel ---
start = timeit.default_timer()
with nvtx.annotate("new cuda.compute kernel", color="green"):
    for _ in range(n_iter):
        with nvtx.annotate("kernel call", color="yellow"):
            run_new(result)
    cupy.cuda.Device().synchronize()
t_new = (timeit.default_timer() - start) / n_iter

total_bytes = input_data.nbytes + offsets.nbytes + result.nbytes
print(f"Segments:         {N_SEGMENTS:,}  x  {SEGMENT_SIZE} elements")
print(f"Input data size:  {total_bytes} bytes ({total_bytes / 1024**2:.1f} MB)")
print(f"Old CUDA kernel:  {t_old:.6f} s/iter")
print(f"unary_transform:  {t_unary:.6f} s/iter")
print(f"New cuda.compute: {t_new:.6f} s/iter")
print(f"Speedup old/unary:{t_old / t_unary:.2f}x")
print(f"Speedup old/new:  {t_old / t_new:.2f}x")

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (awkward3@9333a7e). Learn more about missing BASE report.
⚠️ Report is 36 commits behind head on awkward3.
✅ All tests successful. No failed tests found.

Additional details and impacted files

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