A gpu index to be searched once and reset repeatedly (vs knn_gpu)

[SomeoneSerge]

I'm trying to decide on accommodating approximate 3NN queries of the following form:

while (true) {
  // ...
  // Two collections of vectors, each 1M vectors, all in GPU memory
  X :: 1_000_000 => nChannels => float32
  Y :: 1_000_000 => nChannels => float32
  
  [distances, indices] = knn_gpu(res, X, Y, 3)
  // ...
}

...or, with an index (I assume this can get me further than the brute-force knn_gpu),

index.reserveMemory(1_000_000)

while (true) {
  // ...
  X :: 1_000_000 => nChannels => float16
  Y :: 1_000_000 => nChannels => float16
  if (not index.isTrained or not trainOnce) {
    index.train(Y)
  }
  index.add(Y)
  [distances, indices] = index.search(X, 3)
  index.reset() // Forget Y
  // ...
}

I can afford quite a leeway accuracy-wise if quantization or such is of any help. In fact, I don't care about the distances, only about the ranking. I cannot afford much time on an iteration. I wish to avoid any GPU-CPU roundtrips (AFAIU, OPQ is only implemented for the CPU? Also, I still haven't exactly wrapped my head around what and when is being quantized when I use OPQ/PQ, and how these options interact wit the inverted lists). So far I have tried a simple "IVF24000,PQ4", which costs me 50-200 sec/iteration and I need to go orders of magnitude lower . I noticed that the documentation suggests adding and querying vectors in batches, which I haven't tried, because I was under impression that I would have to be add()ing batches sequentially

My assumption is that all vectors y in Y are sufficiently separated, and that on average every x in X corresponds to a unique y. I expect that if I had allocated a bin per y in Y and put each x in just a handful of nearest bins (maybe using some quantization or hashing, but I'm not an expert in these matters), I would have ended up with more or less equally sized lists that could be traversed in parallel. Am I correct in thinking that under these assumptions, and given that data is already on the GPU, it should be possible to run the above loop at 100s of milliseconds per iteration?

Is faiss an appropriate tool for this task? Would anyone have suggestions on choosing an index and/or preprocessing?

EDIT of 2023-04-07

Alright, my fault, setDefaultNullStreamAllDevices() makes this 20 sec/it, and reserveVecs maybe cuts a few more seconds off the period. Got to get another two orders of magnitude down. I'm still naively using IVF24000,PQ4.

EDIT of 2023-04-14

I see now that reset() only empties the lists, but does not reset the quantizer or remove centroid information.
I also see that re-training is explicitly un-supported: #597.
Looking for a way to re-create an index from scratch without re-allocating the memory or relying on the garbage collector

---

#assignment #bipartite #icp #transport

[namespace-Pt]

index.reset()
index.is_trained = False
quantizer = faiss.downcast_index(index.quantizer)
quantizer.is_trained = False
index.is_trained = False