[FEA] Pass cudaStreamPerThread to numba/CuPy kernels

[shwina]

Memory allocations should already use PTDS since both numba and CuPy allocate memory using RMM. Kernels on the other hand, may explicitly need to be passed the cudaStreamPerThread stream handle.

cc: @jakirkham @kkraus14

[jakirkham]

Some relevant issues listed below:

xref: rapidsai/dask-cuda#96
xref: cupy/cupy#3755
xref: numba/numba#5137

Also some relevant PRs below:

xref: rapidsai/rmm#480
xref: dask/distributed#4034

[github-actions]

This issue has been marked rotten due to no recent activity in the past 90d. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed.

[jakirkham]

cc @pentschev (for awareness)

[vyasr]

@shwina Is this intended to be unconditional, or only when requested by the user? We will also need to start passing streams down to libcudf everywhere in order to support this.

[pentschev]

My take on this is that we could begin with an opt-in solution, like we've done for CuPy with CUPY_CUDA_PER_THREAD_DEFAULT_STREAM. This could allow us to test things without being too intrusive, and if this proves to be useful then we could enable it by default.

[jakirkham]

Yeah agreed. We probably want to support both the current mode and PTDS for some time. There are likely different issues we will run into. So having an escape hatch to old behavior is quite useful.