[BUG] Large memory overhead while subsampling vectors for IVF methods

[tfeher]

Describe the bug

The IVF methods use random subsampling to create a training set for k-means clustering.

The random sampling algorithm allocates several temporary buffers:

raft/cpp/include/raft/random/detail/rng_impl.cuh

Lines 293 to 296 in d4ae271

	rmm::device_uvector<WeightsT> expWts(len, stream);
	rmm::device_uvector<WeightsT> sortedWts(len, stream);
	rmm::device_uvector<IdxT> inIdx(len, stream);
	rmm::device_uvector<IdxT> outIdxBuff(len, stream);

On top of this sortPairs would also allocate temporary buffers, with size O(len)

Here len is the size (number of vectors) in the whole dataset. When we are indexing DEEP-1B, then len=1e9, and the temporary space becomes prohibitive.

Steps/Code to reproduce bug
Run DEEP-1B test with IVF-PQ, with subsample ratio 10. It will run OOM.

Expected behavior
Subsample vectors with minimal memory overhead.

[tfeher]

A quick fix for this problem is to revert the following PRs:

• Add random subsampling for IVF methods #2077 this introduced the memory overhead
• Subsampling for IVF-PQ codebook generation #2052 relies on random sampling for correctness, therefore needs to be reverted if Add random subsampling for IVF methods #2077 is reverted.

I am investigating whether we have an alternative quick fix to reduce the memory overhead for the subsampling.