rapidsai / raft

RAFT contains fundamental widely-used algorithms and primitives for machine learning and information retrieval. The algorithms are CUDA-accelerated and form building blocks for more easily writing high performance applications.
https://docs.rapids.ai/api/raft/stable/
Apache License 2.0
768 stars 194 forks source link

[FEA] IVF index building with pinned H2D transfer #2106

Open tfeher opened 9 months ago

tfeher commented 9 months ago

Is your feature request related to a problem? Please describe. For IVF-Flat ad IVF-PQ index building, large datasets are provided in host memory or as mmap-ed file. After the cluster centers are trained, both method streams through the whole dataset twice. Currently there is no overlap between host to device copies and additional data processing on the GPU.

Describe the solution you'd like Use pinned buffers to copy the data to the GPU and overlap it with GPU side computation.

Additional context

IVF-Flat and IVF-PQ streams through the whole dataset here:

We use batch_load_iterator to copy the data to host. Ideally, we could improve the batch load-iterator to prefetch the data into a pinned buffer.

tfeher commented 9 months ago

Tagging @abc99lr who plans to work on this, and @achirkin for visibility.