[ANN_BENCH] Remove Stream Pool Creation from Individual Wrappers - #2526
[ANN_BENCH] Remove Stream Pool Creation from Individual Wrappers#2526tarang-jain wants to merge 6 commits into
Conversation
|
Could you please give a bit more detail on where the blocking happens? We've been trying to avoid using default streams across raft and cuvs for a long time. If they still appear somewhere within the library, it probably makes sense to fix there? |
achirkin
left a comment
There was a problem hiding this comment.
Thanks, I think the change makes perfect sense and can only improve the performance. We don't need any extra stream syncs, because the replaced managed memory was used as a large workspace resource (which is async by contract).
Just please update the description with the explanation of the problem you wanted to solve and do a small cleanup (see below).
| size_t n_streams = cagra_indices.size(); | ||
| raft::resource::set_cuda_stream_pool(composite_handle, | ||
| std::make_shared<rmm::cuda_stream_pool>(n_streams)); | ||
| make_non_blocking_stream_pool(n_streams)); |
There was a problem hiding this comment.
I'm a bit surprised how this stream pool setting has appeared in all wrapper code in the first place.
Please remove all these repeated set_cuda_stream_pools from the wrapper files and put it exactly once either in the shared_raft_resources or in configured_raft_resources object (where more appropriate).
There was a problem hiding this comment.
I got rid of this behavior in the other ivf wrappers. But for cagra it seems this particular pool is a special case, for the composite index.
| configured_raft_resources() : configured_raft_resources{std::make_shared<shared_raft_resources>()} | ||
| { | ||
| raft::resource::set_cuda_stream_pool( | ||
| *res_, std::make_shared<rmm::cuda_stream_pool>(1, rmm::cuda_stream::flags::non_blocking)); |
There was a problem hiding this comment.
Add a prefetch stream in the default constructor.
…vs-bench-prefetch-async-rmm
non_blocking_streamstruct