check the current device resource - Faiss

RAPIDS Memory Manager (RMM)

We recommend using RMM to configure GPU memory resources. For performance, a best practice with cuVS, as with other libraries in the RAPIDS ecosystem, is to use a pool memory resource to allocate a chunk of memory up front for the current GPU device.

python

import rmm
pool = rmm.mr.PoolMemoryResource(
    rmm.mr.CudaMemoryResource(),
    initial_pool_size=2**30
)
rmm.mr.set_current_device_resource(pool)

python

# check the current device resource
current_resource = rmm.mr.get_current_device_resource()
print(current_resource)

Building a cuVS GPU index

With the faiss-gpu-cuvs package, the cuVS implementation is chosen by default for supported index types and can therefore be used with zero code change. Below contains an example of creating an IVFPQ index using cuVS.

python

import faiss
import numpy as np

np.random.seed(1234)
xb = np.random.random((1000000, 96)).astype('float32')
xq = np.random.random((10000, 96)).astype('float32')
xt = np.random.random((100000, 96)).astype('float32')

res = faiss.StandardGpuResources()
# Disable the default temporary memory allocation since an RMM pool resource has already been set.
res.noTempMemory()

python

# Case 1: Creating cuVS GPU index
config = faiss.GpuIndexIVFPQConfig()
config.interleavedLayout = True
assert(config.use_cuvs)
index_gpu = faiss.GpuIndexIVFPQ(res, 96, 1024, 96, 6, faiss.METRIC_L2, config) # expanded parameter set with cuVS (bits per code = 6).
%time index_gpu.train(xt)
%time index_gpu.add(xb)

python

# Case 2: Cloning a CPU index to a cuVS GPU index
quantizer = faiss.IndexFlatL2(96)
index_cpu = faiss.IndexIVFPQ(quantizer,96, 1024, 96, 8, faiss.METRIC_L2)
index_cpu.train(xt)
co = faiss.GpuClonerOptions()
%time index_gpu = faiss.index_cpu_to_gpu(res, 0, index_cpu, co)

# The cuVS index now uses the trained quantizer as it's IVF centroids.
assert(index_gpu.is_trained)
%time index_gpu.add(xb)
k = 10
%time D, I = index_gpu.search(xq, k)

Build a cuVS CAGRA index

The following example demonstrates building and searching the CAGRA index with FAISS.

python

# Step 1: Create the CAGRA index config
config = faiss.GpuIndexCagraConfig()
config.graph_degree = 32
config.intermediate_graph_degree = 64

# Step 2: Initialize the CAGRA index
res = faiss.StandardGpuResources()
gpu_cagra_index = faiss.GpuIndexCagra(res, 96, faiss.METRIC_L2, config)

# Step 3: Add the 1M vectors to the index
n = 1000000
data = np.random.random((n, 96)).astype('float32')
%time gpu_cagra_index.train(data)

# Step 4: Search the index for top 10 neighbors for each query.
xq = np.random.random((10000, 96)).astype('float32')
%time D, I = gpu_cagra_index.search(xq,10)

CAGRA to HNSW

A CAGRA index can be automatically converted to HNSW through the new faiss.IndexHNSWCagra CPU index class.

python

# Create the HNSW index object.
d = 96
M = 16
cpu_hnsw_index = faiss.IndexHNSWCagra(d, M, faiss.METRIC_L2)
# Create the full HNSW hierarchy
cpu_hnsw_index.base_level_only=False

# Initializes the HNSW base layer with the CAGRA graph. 
%time gpu_cagra_index.copyTo(cpu_hnsw_index)

# Add new vectors to the hierarchy.
newVecs = np.random.random((100000, 96)).astype('float32')
%time cpu_hnsw_index.add(newVecs)