Make the CUDA stream pool per-thread and per-device - #23672
Conversation
The global pool handed overlapping streams to concurrent threads, creating false dependencies between unrelated work. Each thread now owns a pool per device that grows on demand up to LIBCUDF_STREAM_POOL_SIZE, and pools are recycled through a free list when a thread exits so thread churn does not accumulate streams.
I am not sure if this is strictly true, I believe streams only have to be alive while enqueueing work (though not necessarily for the duration of execution). Can you justify this claim? edit: It’s probably because of pool recycling, not stream lifetime. This is fine. |
| streams.emplace_back(_pool.get_stream()); | ||
| grow_to(count); | ||
| auto const first = std::exchange(_next_stream, _next_stream + count); | ||
| auto streams = std::vector<rmm::cuda_stream_view>(); |
There was a problem hiding this comment.
Include <cuda/stream> and replace all uses of rmm::cuda_stream_view with cuda::stream_ref.
The interface documented every accessor as thread safe, which the unsynchronized per-thread implementation does not provide; state instead that a pool is owned by one thread at a time. The tests derived their expectations from an assumed pool size, so they depended on which recycled pool a thread happened to adopt and on the configured cap; they now derive sizes from the pool and skip where a property is undefined. Adds coverage for the per-device pools.
|
Be aware that there are hardware limits to how many actual parallel workstreams a program can use at once and that if you exceed that you will end up with the same implicit synchronization happening even if you produce these additional software level streams. Your change is still worthwhile to provide us some extra concurrency, just want to point out that it will top out. There is a degree of control of the underlying hardware limits available with environment variables. |
The pool interface now hands out cuda::stream_ref instead of rmm::cuda_stream_view, and owns cuda::stream rather than rmm::cuda_stream. cuda::stream is always non-blocking and takes the device explicitly, which suits a per-device pool, and its destructor uses the driver API so it does not depend on the current device. Callers are unaffected where they pass streams on, since the two view types convert implicitly; the changes elsewhere are value() to get() at kernel launches and synchronize() to sync().
Drops the growth high-water-mark and nested-fork tests, which asserted a growth schedule and a rotation heuristic the interface does not guarantee, and the per-device test, which always skips because tests request a single GPU. What remains covers cross-thread disjointness, pool reuse after a thread exits, and the over-cap repeat contract; a break in any of those is otherwise silent.
Description
libcudf's stream pool was a single global pool of 32 streams shared by every thread, so concurrent threads calling
fork_streamswere handed overlapping streams. That creates false dependencies between unrelated work and serializes it. Each thread now owns a pool per device, created empty and grown on demand up to a cap configurable withLIBCUDF_STREAM_POOL_SIZE(default 32). A request grows the pool to twice the requested count and hands out streams starting at a rotating offset, so a nestedfork_streamsgenerally avoids the streams its caller already holds. Pools are returned to a per-device free list when a thread exits and adopted by the next thread that needs one, so applications that create and destroy many threads do not accumulate streams.fork_streamsandjoin_streamskeep their signatures. Streams may still be used from other threads, which libcudf does in the host compression, CSV, and JSON read paths, but they must not outlive the thread that acquired them; all current call sites join before the acquiring frame returns.Multithreaded parquet reads improve by 7-12% in GPU time at 8 threads (
parquet_multithreaded_read_decode_mixed27.4 ms to 25.5 ms,parquet_multithreaded_read_decode_string29.5 ms to 26.8 ms), with run-to-run noise dropping from 9-11% to 4-6%. Single- and two-thread configs, ORC, and groupby are unchanged, and peak memory is unchanged.Checklist