SNMG ANN

[viclafargue]

The goal of this PR is to implement a distributed (single-node-multiple-GPUs) implementation of ANN indexes. It will allow to build, extend and search an index on multiple GPUs.

Before building the index, the user has to choose between two modes :

• Sharding mode : The index dataset is split, each GPU trains its own index with its respective share of the dataset. This is intended to both increase the search throughput and the maximal size of the index.
• Index duplication mode : The index is built once on a GPU and then copied over to others. Alternatively, the index dataset is sent to each GPU to be built there. This intended to increase the search throughput.

[copy-pr-bot[bot]]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[viclafargue]

The PR is ready for a first review. In its current state, it implements the build, extend and search ANN methods (IVF-Flat and IVF-PQ only for now) in index duplication and sharding mode. For now, the index duplication mode only works by copying the index dataset over and building the index on each GPU separately. I am now looking to improve the API in such a way that it would allow to build the index on a GPU and copy it over. Serialization on disk would work, but does not seem ideal. Then, transferring the index attributes through NCCL seem like not very safe. What would you recommend?

[lijinf2]

Thank you Victor. I have learned a lot from the code! I like the idea of combining three algorithms into one unified interface. A few questions to make myself more familiar with the PR and design choice. Will be wonderful if Spark Rapids ML can leverage the APIs in this PR.

[review-notebook-app[bot]]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[cjnolet]

/ok to test