Optimize transfer/compute overlap in out-of-core KMeans - #2538
Optimize transfer/compute overlap in out-of-core KMeans#2538viclafargue wants to merge 4 commits into
Conversation
viclafargue
added
improvement
|
/ok to test 4b8a5c5 |
| raft::resource::get_stream_pool_size(handle) >= 1) { | ||
| batch_copy_stream = raft::resource::get_stream_from_stream_pool(handle); | ||
| } else { | ||
| owned_batch_copy_stream.emplace(rmm::cuda_stream::flags::non_blocking); |
There was a problem hiding this comment.
Lets not create a prefetch stream on our own. It is solely controlled by the raft resource. The resource should give full control of concurrent streams to the user.
| IndexT row_width, | ||
| IndexT batch_size, | ||
| rmm::cuda_stream_view copy_stream, | ||
| rmm::device_async_resource_ref mr) |
There was a problem hiding this comment.
hardcoding rmm::device_async_resource_ref wont give us problems? What if the currently set MR is not an async resource?
| { | ||
| auto const& batch = batches_[pos]; | ||
| const auto bytes = batch.size * row_width_ * sizeof(DataT); | ||
| RAFT_CUDA_TRY(cudaMemcpyAsync( |
There was a problem hiding this comment.
Correct me if I am wrong but the main new addition here in comparison to the ANN batch loader is the recycling at the final slots, right? So prefetch continues for the next iteration. Having an entirely new batch loader makes me a bit uneasy. Prefetch should not be algorithm specific. Can we combine this with the existing ANN batch loader (this can be done in a follow up)?
2 participants
Summary
This PR improves data-transfer and compute overlap for host-resident out-of-core KMeans using a cyclic two-buffer pipeline.
Implementation
Benchmark under similar configuration
10 GiB pinned-host FP32 dataset (10,485,760 × 256), 2,560 clusters, three iterations, ten 1 GiB out-of-core batches, and 131,072-sample assignment tiles.
mainProfile
Main branch :
This PR :
This PR (multi-GPU) :
