[FEA]: Introduce cuda.cooperative overloads not requiring temporary storage

[gevtushenko]

Is this a duplicate?

• I confirmed there appear to be no duplicate issues for this request and that I agree to the Code of Conduct

Area

CUB

Is your feature request related to a problem? Please describe.

cuda.cooperative API currently has an issue. We do not specify alignment of the temporary storage. This leads to bugs like the following one:

cx_int_sum = cudax.block.sum(numba.int32, 256)
int_sum_storage_bytes = cx_int_sum.temp_storage_bytes
cx_float_sum = cudax.block.sum(dtype=numba.float64, threads_in_block=256)
float_sum_storage_bytes = cx_float_sum.temp_storage_bytes

cx_files = cx_float_sum.files + cx_int_sum.files


@numba.cuda.jit(
    "void(int32[:], float64[:])",
    link=cx_files
)
def kernel(args_in, args_out):
    int_sum_storage = numba.cuda.shared.array(shape=int_sum_storage_bytes, dtype=numba.uint8)
    float_sum_storage = numba.cuda.shared.array(shape=float_sum_storage_bytes, dtype=numba.uint8)
    # do int reduction first
    tix = numba.cuda.threadIdx.x
    val = cx_int_sum(int_sum_storage, args_in[tix])
    val2 = cx_float_sum(float_sum_storage, numba.float64(args_in[tix]))
    args_out[tix] = val + val2

Because both allocations of shared memory are made at uint8 granularity, second one is not properly aligned, leading to:

cupy_backends.cuda.api.runtime.CUDARuntimeError: cudaErrorMisalignedAddress: misaligned address

Describe the solution you'd like

Majority of kernels do not create temporary storage unions, so we could simplify the API by not requiring temporary storage:

    block_reduce = cudax.block.reduce(numba.int32, threads_in_block, op)

    @cuda.jit(link=block_reduce.files)
    def kernel(input, output):
        block_output = block_reduce(input[cuda.threadIdx.x])

Describe alternatives you've considered

No response

Additional context

No response