[FEA] Multiple threads sharing the same GPU

[revans2]

Is your feature request related to a problem? Please describe.
Currently all of cudf processing uses a single stream which blocks all other streams from sharing the GPU. The plugin uses a semaphore to prevent too many processes on the GPU at once. This works to overlap some I/O with processing, but it does not work if you want to overlap processing.

Describe the solution you'd like
I would be great to be able to support multiple streams, possibly using per-thread-default streams.

[rongou]

The consensus from RMM and cuDF is to rely on per-thread default stream (PTDS).

[rongou]

Here is the design proposal.

Motivation

Spark is multi-threaded on the host (CPU) side. Input is divided into partitions, and multiple tasks are launched on each CPU to work on these partitions. Each task typically runs in a single thread, obtained from a thread pool. This task level parallelism ensures all the CPU cores are utilized to increase overall throughput.

With the spark-rapids plugin, some or all of the work in each task is offloaded to the GPU. To efficiently use the GPU, the CUDA operations in these tasks should execute concurrently. CUDA applications manage concurrency by executing asynchronous commands in streams, sequences of commands that execute in order. Currently all host threads share the single default CUDA stream. Also known as the legacy default stream or the NULL stream, it implicitly synchronizes with itself and all other streams on the device, causing CUDA operations to be serialized.

To increase GPU throughput, it is desirable to execute each task or host thread in a separate, asynchronous stream. Ideally the legacy default stream is never used so as not to trigger implicit synchronization.

Goals

• Each Spark task or host thread runs in a separate CUDA stream.
• The legacy default stream is not used.
• Test and fix the current code base so that concurrent execution on the GPU works correctly, e.g. no race conditions, deadlocks etc.

Non-Goals

• Once tasks execute concurrently on GPUs, there may be further opportunities to tune current approaches to get better performance. This is outside the scope of this issue.
• Runtime compilation using NVRTC/Jitify is only for device code, which is oblivious to streams (unless using CUDA dynamic parallelism).

Assumptions

• There shouldn’t be any big structural changes to the code base, either the underlying C++ libraries, or the Java/Scala code in the spark-rapids plugin.
• A Spark task shouldn’t need to manipulate CUDA streams directly, or deal with multiple streams.
• Some background tasks such as asynchronous spill may operate with their own streams.
• GPU memories allocated in one task/stream may be deallocated in another task/stream.

Risks

• Concurrent execution of GPU code may be hard to reason about or debug. We need good test coverage to make sure the code works correctly.
• Some third party libraries may never be compiled with per-thread default stream enabled. In the case of XGBoost or other machine learning libraries, since model training does not overlap with ETL, using the legacy default stream during training should not have an adverse effect. However, libraries used during ETL may cause excessive synchronization when PTDS is not enabled.
• There are upcoming CUDA features that may impact some key components in the stack.

Design

CUDA 7 introduced the per-thread default stream (PTDS), which has two effects. First, it gives each host thread its own default stream. This means that commands issued to the default stream by different host threads can run concurrently. Second, these default streams are regular streams. This means that commands in the default stream may run concurrently with commands in non-default streams.

To enable PTDS, we can compile with the nvcc command-line option --default-stream per-thread. Once this is done on all CUDA code, each Spark thread would use its own default stream, and operations from different threads would execute concurrently on the GPU.

The work involved can be roughly divided into two parts, basic mechanism and correctness.

Basic Mechanism

• Option to enable PTDS for RMM ([REVIEW] add option for per-thread default stream rapidsai/rmm#354)
• Option to enable PTDS for cuDF ([REVIEW] Add CMake option for per-thread default stream rapidsai/cudf#4995)
• Option to enable PTDS for cuDF JNI ([REVIEW] Add CMake option for per-thread default stream rapidsai/cudf#4995)
• Ensure the Thrust library, which is heavily used by cuDF, handles PTDS correctly (support per-thread default stream thrust#1128)
• RMM should use the correct version of Thrust ([REVIEW] fetch thrust/cub from github rapidsai/rmm#378)
• cuDF should use the correct version of Thrust ([REVIEW] fetch thrust/cub from github rapidsai/cudf#5315)
• cuDF JNI should use the correct version of Thrust ([REVIEW] use new release of thrust/cub in jni build [skip ci] rapidsai/cudf#5414)

Correctness

• Test PTDS in Thrust (https://github.com/thrust/thrust/issues/1132, Add tests for legacy and per-thread default streams thrust#1201)
• Support PTDS in RMM ([FEA] Add per-thread default stream support to pool_memory_resource rapidsai/rmm#410, [REVIEW] Add per-thread-default stream support to pool_memory_resource using thread-local CUDA events rapidsai/rmm#425, [REVIEW] switch from CNMeM to pool_memory_resource [skip ci] rapidsai/cudf#5632)
• Test PTDS in RMM/cuDF ([DISCUSSION] What additional libcuDF and cuDF unit testing is needed for Per-Thread Default Stream rapidsai/cudf#5614, [REVIEW] enable per-thread default stream in gpu ci build rapidsai/rmm#436, [REVIEW] enable per-thread default stream in gpu ci rapidsai/cudf#5673)
• Test PTDS in cuDF JNI

Finally, we should run some benchmarks with PTDS enabled. These can be standard TPC-H or TPC-DS queries, or real customer queries if possible.

• Benchmarks

After verifying that everything works correctly, the spark-rapids distribution of RMM/cuDF can enable PTDS by default.

• Enable PTDS in spark-rapids distribution

Alternatives Considered

Most (if not all) cuDF functions support passing in an extra parameter for the stream to execute on, although currently these stream parameters are not exposed in cuDF’s public API. One alternative is to make these parameters public and for the spark-rapids plugin to manage CUDA streams explicitly. Each task can grab its own stream, preferably from a stream pool, and pass it to all subsequent cuDF calls. The benefit of this approach is that the intent is more explicit, and the code doesn’t need to be compiled with the per-thread default stream flag. The drawback is that the client code needs to manage streams, leading to more complexity.

[jrhemstad]

• Some third party libraries may never be compiled with per-thread default stream enabled.

fyi, you can explicitly pass cudaStreamPerThread into external libraries to get PTDS behavior even if the external library wasn't complied with PTDS enabled.

[rongou]

Closing this as the list of items is done. Will open another one for performance tuning.