[EPIC] Investigate refactoring CuPy to use cuda.parallel

[jrhemstad]

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

cuda.parallel (Python)

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

Today, CuPy uses Thrust/CUB algorithms to implement much of it's functionality. That works today by precompiling Thrust algorithms for a variety of fixed types. This is undesirable for a few reasons: it increases binary size, it limits exposure of some algorithms (like segmented sort) due to combinatorial type explosion.

cuda.parallel can and should be able to replace any existing use of pre-instantiated Thrust/CUB algorithms and provide a few benefits:

• Reduce binary size (going to JIT)
• Custom type support
• Custom operator support
• Additional algorithm support (because JIT avoids the type combination problem)

Describe the solution you'd like

To start, we'd like to have an inventory of what Thrust/CUB stuff CuPy is using today and where.

From there, we should investigate how we can use cuda.parallel.reduce_into to replace existing uses of cub::DeviceReduce/thrust::reduce.

Describe alternatives you've considered

No response

Additional context

No response

[shwina]

To give an idea of the binary size improvement this could bring, here are the sizes of the largest 10 extensions in the CuPy build:

~/workspace/cupy$ find . -type f -name "*.so" -exec ls -lhS {} + | awk '{print $5, $9}' | head -20 | grep -v "./build"

42M ./cupy/cuda/thrust.cpython-310-x86_64-linux-gnu.so
34M ./cupy/cuda/cub.cpython-310-x86_64-linux-gnu.so
7.3M ./cupy/random/_generator_api.cpython-310-x86_64-linux-gnu.so
1.5M ./cupy_backends/cuda/libs/cusparse.cpython-310-x86_64-linux-gnu.so
1.4M ./cupy_backends/cuda/api/runtime.cpython-310-x86_64-linux-gnu.so
1.2M ./cupy/cuda/jitify.cpython-310-x86_64-linux-gnu.so
1.2M ./cupy_backends/cuda/libs/cusolver.cpython-310-x86_64-linux-gnu.so
851K ./cupy/_core/core.cpython-310-x86_64-linux-gnu.so
693K ./cupy/_core/fusion.cpython-310-x86_64-linux-gnu.so
671K ./cupy/cuda/memory.cpython-310-x86_64-linux-gnu.so

cupy.cuda.thrust and cupy.cuda.cub are by far the largest extension modules.