Add CuVS and use for GPU Clustering

CMakeLists.txt

@@ -12,18 +12,6 @@ set(Torch_USE_CUDA OFF CACHE BOOL "Force disable CUDA in Torch")
 set(Torch_NO_CUDA ON CACHE BOOL "Force disable CUDA in Torch")
 set(USE_CUDA OFF CACHE BOOL "Force disable CUDA globally")
 
-# QUAKE_ENABLE_GPU: Enable GPU support for Faiss
-# Default: OFF
-if(QUAKE_ENABLE_GPU)
-    set(FAISS_ENABLE_GPU ON)
-else()
-    set(FAISS_ENABLE_GPU OFF)
-endif()
-
-if(QUAKE_ENABLE_GPU)
-    add_compile_definitions(FAISS_ENABLE_GPU)
-endif()
-
 if(QUAKE_USE_NUMA)
     add_compile_definitions(QUAKE_USE_NUMA)
 endif()
@@ -64,6 +52,25 @@ set(project_BINDINGS_DIR ${CPP_SOURCE}/bindings)
 set(project_THIRD_PARTY_DIR ${CPP_SOURCE}/third_party)
 set(project_TEST_DIR test/cpp)
 
+if(QUAKE_ENABLE_GPU)
+    find_package(CuVS REQUIRED)
+    include(${project_THIRD_PARTY_DIR}/cmake/fetch_rapids.cmake)
+    include(rapids-cmake)
+    include(rapids-cpm)
+    include(rapids-cuda)
+    include(rapids-export)
+    include(rapids-find)
+
+    rapids_cuda_init_architectures(quake_c)
+
+    rapids_cpm_init()
+    set(BUILD_CUVS_C_LIBRARY OFF)
+    include(${project_THIRD_PARTY_DIR}/cmake/get_cuvs.cmake)
+
+    add_compile_definitions(QUAKE_ENABLE_GPU)
+endif()
+
+set(FAISS_ENABLE_GPU OFF)
 # ---------------------------------------------------------------
 # Print out Compiler and Path Information
 # ---------------------------------------------------------------
@@ -81,7 +88,6 @@ message(STATUS "QUAKE_USE_NUMA: ${QUAKE_USE_NUMA}")
 # Apple-specific adjustments
 if(APPLE)
     include_directories("/opt/homebrew/opt/openblas/include")
-    set(FAISS_ENABLE_GPU OFF)
 endif()
 
 # Compiler options and definitions
@@ -149,6 +155,10 @@ elseif(UNIX)
     if(QUAKE_USE_NUMA)
         list(APPEND LINK_LIBS -lnuma)
     endif()
+
+    if (QUAKE_ENABLE_GPU)
+        list(APPEND LINK_LIBS cuvs::cuvs)
+    endif()
 else()
     # unsupported platform
     message(FATAL_ERROR "Unsupported platform")
@@ -211,7 +221,7 @@ endif()
 # ---------------------------------------------------------------
 message(STATUS "--------- Final Configuration Summary ---------")
 message(STATUS "Build type: ${CMAKE_BUILD_TYPE}")
-message(STATUS "GPU Enabled: ${FAISS_ENABLE_GPU}")
+message(STATUS "GPU Enabled: ${QUAKE_ENABLE_GPU}")
 message(STATUS "NUMA Enabled: ${QUAKE_USE_NUMA}")
 message(STATUS "Python used: ${Python3_EXECUTABLE}")
 message(STATUS "Torch Path: ${TorchPath}")
@@ -221,4 +231,4 @@ message(STATUS "MKL_LINK: ${MKL_LINK}")
 message(STATUS "MKL_INTERFACE_FULL: ${MKL_INTERFACE_FULL}")
 message(STATUS "MKL_THREADING: ${MKL_THREADING}")
 message(STATUS "MKL_MPI: ${MKL_MPI}")
-message(STATUS "------------------------------------------------")
+message(STATUS "------------------------------------------------")

environments/ubuntu-cuda/conda.yaml → environments/ubuntu-cuda118/conda.yaml

@@ -3,13 +3,17 @@ channels:
   - pytorch
   - defaults
   - conda-forge
+  - nvidia
+  - libcuvs
 dependencies:
   - python=3.11
   - numpy
   - pandas
   - faiss-gpu
   - matplotlib
   - pytest
+  - libcuvs
+  - cuda-version=12.8
   - pip
   - pip:
     - sphinx

environments/ubuntu-cuda124/Dockerfile

@@ -0,0 +1,64 @@
+# Use a CUDA-enabled Ubuntu base image
+FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
+
+# -----------------------------
+# Set up environment variables
+# -----------------------------
+ENV CONDA_DIR=/opt/miniconda
+ENV PATH="${CONDA_DIR}/bin:${PATH}"
+ENV DEBIAN_FRONTEND=noninteractive
+
+# -----------------------------
+# Install system dependencies
+# -----------------------------
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+      wget \
+      curl \
+      build-essential \
+      ca-certificates \
+      swig \
+      git \
+      libomp5 \
+      libomp-dev \
+      graphviz \
+      libnuma-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+
+# Install CMake 3.24.2
+RUN wget -qO /tmp/cmake.sh https://github.com/Kitware/CMake/releases/download/v3.30.4/cmake-3.30.4-linux-x86_64.sh && \
+    chmod +x /tmp/cmake.sh && \
+    /tmp/cmake.sh --skip-license --prefix=/usr/local && \
+    rm /tmp/cmake.sh
+
+
+# -----------------------------
+# Install Miniconda
+# -----------------------------
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh && \
+    bash /tmp/miniconda.sh -b -p $CONDA_DIR && \
+    rm /tmp/miniconda.sh
+
+# -----------------------------
+# Copy in your conda environment YAML
+# -----------------------------
+COPY environments/ubuntu-cuda/conda.yaml /tmp/conda.yaml
+
+# Create quake-env
+RUN conda env create -f /tmp/conda.yaml && conda clean -afy
+
+# -----------------------------
+# Install GPU-enabled PyTorch
+# -----------------------------
+RUN conda run -n quake-env pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu124
+
+# -----------------------------
+# Debug: show conda information
+# -----------------------------
+RUN echo "===== DEBUG: which conda =====" && which conda
+RUN echo "===== DEBUG: conda info =====" && conda info
+RUN echo "===== DEBUG: conda env list =====" && conda env list
+RUN echo "===== DEBUG: quake-env check =====" && conda run -n quake-env python -c "import sys; print('OK in quake-env; python:', sys.executable)"
+
+CMD ["/bin/bash"]

environments/ubuntu-cuda124/conda.yaml

@@ -0,0 +1,24 @@
+name: quake-env
+channels:
+  - rapidsai
+  - pytorch
+  - defaults
+  - conda-forge
+  - nvidia
+dependencies:
+  - python=3.11
+  - numpy
+  - pandas
+  - faiss-gpu
+  - matplotlib
+  - pytest
+  - libcuvs
+  - cuda-nvrtc-dev=12.4
+  - cuda-version=12.4
+  - pip
+  - pip:
+    - sphinx
+    - sphinx_rtd_theme
+    - sphinxcontrib-mermaid
+    - graphviz
+    - pyyaml

src/cpp/include/clustering.h

@@ -11,10 +11,41 @@
 
 class IndexPartition;
 
+/**
+ * @brief Clusters vectors into partitions using faiss::Clustering
+ *
+ *
+ * @param vectors The vectors to cluster.
+ * @param ids The IDs of the vectors.
+ * @param n_clusters The number of clusters to create.
+ * @param metric_type The metric type to use for clustering.
+ * @param niter The number of iterations to run k-means.
+ * @param initial_centroids The initial centroids to use for k-means.
+ */
+shared_ptr<Clustering> kmeans_cpu(Tensor vectors,
+                              Tensor ids,
+                              shared_ptr<IndexBuildParams> build_params,
+                              Tensor initial_centroids = Tensor());
+
+/**
+ * @brief Clusters vectors into partitions using CuVS k-means.
+ *
+ *
+ * @param vectors The vectors to cluster.
+ * @param ids The IDs of the vectors.
+ * @param n_clusters The number of clusters to create.
+ * @param metric_type The metric type to use for clustering.
+ * @param niter The number of iterations to run k-means.
+ * @param initial_centroids The initial centroids to use for k-means.
+ */
+#ifdef QUAKE_ENABLE_GPU
+shared_ptr<Clustering> kmeans_cuvs_sample_and_predict(
+    Tensor vectors, Tensor ids, shared_ptr<IndexBuildParams> build_params);
+#endif
+
 /**
  * @brief Clusters vectors into partitions using k-means.
  *
- * Uses the faiss::Clustering class to cluster vectors into n_clusters partitions.
  *
  * @param vectors The vectors to cluster.
  * @param ids The IDs of the vectors.
@@ -25,10 +56,7 @@ class IndexPartition;
  */
 shared_ptr<Clustering> kmeans(Tensor vectors,
                               Tensor ids,
-                              int n_clusters,
-                              MetricType metric_type,
-                              int niter = 5,
-                              bool use_gpu = false,
+                              shared_ptr<IndexBuildParams> build_params,
                               Tensor initial_centroids = Tensor());
 
 

src/cpp/include/common.h

@@ -38,14 +38,6 @@
 #include <numaif.h>
 #endif
 
-#ifdef FAISS_ENABLE_GPU
-#include <faiss/gpu/GpuIndexFlat.h>
-#include <faiss/gpu/GpuIndexIVFFlat.h>
-#include <faiss/gpu/GpuIndexIVFPQ.h>
-#include <faiss/gpu/StandardGpuResources.h>
-#include <faiss/gpu/GpuCloner.h>
-#endif
-
 using torch::Tensor;
 using std::vector;
 using std::unordered_map;
@@ -71,6 +63,8 @@ constexpr int DEFAULT_NLIST = 0;                   ///< Default number of cluste
 constexpr int DEFAULT_NITER = 5;                   ///< Default number of k-means iterations used during clustering.
 constexpr const char* DEFAULT_METRIC = "l2";       ///< Default distance metric (either "l2" for Euclidean or "ip" for inner product).
 constexpr int DEFAULT_NUM_WORKERS = 0;             ///< Default number of workers (0 means single-threaded).
+constexpr int DEFAULT_GPU_BATCH_SIZE = 100000;             ///< Default batch size for GPU index building.
+constexpr int DEFAULT_GPU_SAMPLE_SIZE = 1000000;           ///< Default sample size for GPU index building.
 
 // Default constants for search parameters
 constexpr int DEFAULT_K = 1;                             ///< Default number of neighbors to return.
@@ -132,11 +126,15 @@ struct IndexBuildParams {
 
     bool use_adaptive_nprobe = false;
     bool use_numa = false;
-    bool use_gpu = false;
     bool verify_numa = false;
     bool same_core = true;
     bool verbose = false;
 
+    // gpu index build params
+    bool use_gpu = false;
+    int gpu_batch_size = DEFAULT_GPU_BATCH_SIZE;
+    int gpu_sample_size = DEFAULT_GPU_SAMPLE_SIZE;
+
     shared_ptr<IndexBuildParams> parent_params = nullptr;
 
     IndexBuildParams() = default;