IndexIVFPQ search first pass very slow

[Kh4L]

Summary

First pass of IndexIVFPQ search takes about 60s, even with a fairly small dataset.

1st search time 60.47380566596985
2nd search time 0.0010259151458740234

Platform

Faiss version: https://github.com/facebookresearch/faiss/tree/v1.7.4

Installed from: source

Faiss compilation options: cmake -B build . && make -C build -j faiss && make -C build -j swigfaiss

Running on:

• CPU
• GPU

Interface:

• C++
• Python

Reproduction instructions

import torch
import faiss
import faiss.contrib.torch_utils
import time

dataset = torch.rand((2**14, 128)).cuda()

query_k = 12
num_cells = 10
num_cells_to_visit = 10
bits_per_vector = 8
channels = dataset.size(1)

index = faiss.IndexIVFPQ(
    faiss.IndexFlatL2(channels),
    channels,
    num_cells,
    bits_per_vector,
    8,
    faiss.METRIC_L2,
)
index.nprobe = num_cells_to_visit

index = faiss.index_cpu_to_gpu(
    faiss.StandardGpuResources(),
    dataset.device.index,
    index,
)

index.train(dataset)
index.add(dataset.detach())

emb = torch.rand((512,128)).cuda()
start = time.time()
score, idx = index.search(emb.detach(), query_k)
print(f"1st search time {time.time() - start}")
emb = torch.rand((512,128)).cuda()
start = time.time()
score, idx = index.search(emb.detach(), query_k)
print(f"2nd search time {time.time() - start}")

[Kh4L]

Profile with nsys, `pass1SelectLists` and `pass2SelectLists` are very slow and it's mostly CPU

[mdouze]

It is probably compiling some kernels on the first run. Do you have a NVIDIA cache directory? Normally in ~/.nv

[Kh4L]

It is probably compiling some kernels on the first run. Do you have a NVIDIA cache directory? Normally in ~/.nv

I don't see any ~/.nv after the first run, and the first pass is consistently very long (60+ s) across runs.

[sunxiaojie99]

I also encountered this problem