Add Cagra-Q compression to the python and rust api's (rapidsai#68)

Authors:
  - Ben Frederickson (https://github.com/benfred)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: rapidsai#68

python/cuvs/cuvs/neighbors/cagra/__init__.py

@@ -13,6 +13,20 @@
 # limitations under the License.
 
 
-from .cagra import Index, IndexParams, SearchParams, build_index, search
+from .cagra import (
+    CompressionParams,
+    Index,
+    IndexParams,
+    SearchParams,
+    build_index,
+    search,
+)
 
-__all__ = ["Index", "IndexParams", "SearchParams", "build_index", "search"]
+__all__ = [
+    "CompressionParams",
+    "Index",
+    "IndexParams",
+    "SearchParams",
+    "build_index",
+    "search",
+]

python/cuvs/cuvs/neighbors/cagra/cagra.pxd

@@ -34,11 +34,22 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil:
         IVF_PQ
         NN_DESCENT
 
+    ctypedef struct cuvsCagraCompressionParams:
+        uint32_t pq_bits
+        uint32_t pq_dim
+        uint32_t vq_n_centers
+        uint32_t kmeans_n_iters
+        double vq_kmeans_trainset_fraction
+        double pq_kmeans_trainset_fraction
+
+    ctypedef cuvsCagraCompressionParams* cuvsCagraCompressionParams_t
+
     ctypedef struct cuvsCagraIndexParams:
         size_t intermediate_graph_degree
         size_t graph_degree
         cuvsCagraGraphBuildAlgo build_algo
         size_t nn_descent_niter
+        cuvsCagraCompressionParams_t compression
 
     ctypedef cuvsCagraIndexParams* cuvsCagraIndexParams_t
 
@@ -74,6 +85,12 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil:
 
     ctypedef cuvsCagraIndex* cuvsCagraIndex_t
 
+    cuvsError_t cuvsCagraCompressionParamsCreate(
+        cuvsCagraCompressionParams_t* params)
+
+    cuvsError_t cuvsCagraCompressionParamsDestroy(
+        cuvsCagraCompressionParams_t index)
+
     cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params)
 
     cuvsError_t cuvsCagraIndexParamsDestroy(cuvsCagraIndexParams_t index)

python/cuvs/cuvs/neighbors/cagra/cagra.pyx

@@ -43,6 +43,82 @@ from libc.stdint cimport (
 from cuvs.common.exceptions import check_cuvs
 
 
+cdef class CompressionParams:
+    """
+    Parameters for VPQ Compression
+
+    Parameters
+    ----------
+    pq_bits: int
+        The bit length of the vector element after compression by PQ.
+        Possible values: [4, 5, 6, 7, 8]. The smaller the 'pq_bits', the
+        smaller the index size and the better the search performance, but
+        the lower the recall.
+    pq_dim: int
+        The dimensionality of the vector after compression by PQ. When zero,
+        an optimal value is selected using a heuristic.
+    vq_n_centers: int
+        Vector Quantization (VQ) codebook size - number of "coarse cluster
+        centers". When zero, an optimal value is selected using a heuristic.
+    kmeans_n_iters: int
+        The number of iterations searching for kmeans centers (both VQ & PQ
+        phases).
+    vq_kmeans_trainset_fraction: float
+        The fraction of data to use during iterative kmeans building (VQ
+        phase). When zero, an optimal value is selected using a heuristic.
+    vq_kmeans_trainset_fraction: float
+        The fraction of data to use during iterative kmeans building (PQ
+        phase). When zero, an optimal value is selected using a heuristic.
+    """
+    cdef cuvsCagraCompressionParams * params
+
+    def __cinit__(self):
+        check_cuvs(cuvsCagraCompressionParamsCreate(&self.params))
+
+    def __dealloc__(self):
+        check_cuvs(cuvsCagraCompressionParamsDestroy(self.params))
+
+    def __init__(self, *,
+                 pq_bits=8,
+                 pq_dim=0,
+                 vq_n_centers=0,
+                 kmeans_n_iters=25,
+                 vq_kmeans_trainset_fraction=0.0,
+                 pq_kmeans_trainset_fraction=0.0):
+        self.params.pq_bits = pq_bits
+        self.params.pq_dim = pq_dim
+        self.params.vq_n_centers = vq_n_centers
+        self.params.kmeans_n_iters = kmeans_n_iters
+        self.params.vq_kmeans_trainset_fraction = vq_kmeans_trainset_fraction
+        self.params.pq_kmeans_trainset_fraction = pq_kmeans_trainset_fraction
+
+    @property
+    def pq_bits(self):
+        return self.params.pq_bits
+
+    @property
+    def pq_dim(self):
+        return self.params.pq_dim
+
+    @property
+    def vq_n_centers(self):
+        return self.params.vq_n_centers
+
+    @property
+    def kmeans_n_iters(self):
+        return self.params.kmeans_n_iters
+
+    @property
+    def vq_kmeans_trainset_fraction(self):
+        return self.params.vq_kmeans_trainset_fraction
+
+    @property
+    def pq_kmeans_trainset_fraction(self):
+        return self.params.pq_kmeans_trainset_fraction
+
+    def get_handle(self):
+        return <size_t>self.params
+
 cdef class IndexParams:
     """
     Parameters to build index for CAGRA nearest neighbor search
@@ -64,17 +140,30 @@ cdef class IndexParams:
             - nn_descent (experimental) will use the NN-Descent algorithm for
               building the knn graph. It is expected to be generally
               faster than ivf_pq.
+    compression: CompressionParams, optional
+        If compression is desired should be a CompressionParams object. If None
+        compression will be disabled.
     """
 
     cdef cuvsCagraIndexParams* params
 
+    # hold on to a reference to the compression, to keep from being GC'ed
+    cdef public object compression
+
+    def __cinit__(self):
+        check_cuvs(cuvsCagraIndexParamsCreate(&self.params))
+        self.compression = None
+
+    def __dealloc__(self):
+        check_cuvs(cuvsCagraIndexParamsDestroy(self.params))
+
     def __init__(self, *,
                  metric="sqeuclidean",
                  intermediate_graph_degree=128,
                  graph_degree=64,
                  build_algo="ivf_pq",
-                 nn_descent_niter=20):
-        cuvsCagraIndexParamsCreate(&self.params)
+                 nn_descent_niter=20,
+                 compression=None):
 
         # todo (dgd): enable once other metrics are present
         # and exposed in cuVS C API
@@ -87,6 +176,10 @@ cdef class IndexParams:
         elif build_algo == "nn_descent":
             self.params.build_algo = cuvsCagraGraphBuildAlgo.NN_DESCENT
         self.params.nn_descent_niter = nn_descent_niter
+        if compression is not None:
+            self.compression = compression
+            self.params.compression = \
+                <cuvsCagraCompressionParams_t><size_t>compression.get_handle()
 
     # @property
     # def metric(self):

python/cuvs/cuvs/test/test_cagra.py

@@ -38,6 +38,7 @@ def run_cagra_build_search_test(
     inplace=True,
     add_data_on_build=True,
     search_params={},
+    compression=None,
 ):
     dataset = generate_data((n_rows, n_cols), dtype)
     if metric == "inner_product":
@@ -49,6 +50,7 @@ def run_cagra_build_search_test(
         intermediate_graph_degree=intermediate_graph_degree,
         graph_degree=graph_degree,
         build_algo=build_algo,
+        compression=compression,
     )
 
     if array_type == "device":
@@ -173,3 +175,11 @@ def test_cagra_index_params(params):
         compare=False,
         build_algo=params["build_algo"],
     )
+
+
+def test_cagra_vpq_compression():
+    dim = 64
+    pq_len = 2
+    run_cagra_build_search_test(
+        n_cols=dim, compression=cagra.CompressionParams(pq_dim=dim / pq_len)
+    )

rust/cuvs/src/cagra/index.rs

@@ -106,8 +106,7 @@ mod tests {
     use ndarray_rand::rand_distr::Uniform;
     use ndarray_rand::RandomExt;
 
-    #[test]
-    fn test_cagra_index() {
+    fn test_cagra(build_params: IndexParams) {
         let res = Resources::new().unwrap();
 
         // Create a new random dataset to index
@@ -117,7 +116,6 @@ mod tests {
             ndarray::Array::<f32, _>::random((n_datapoints, n_features), Uniform::new(0., 1.0));
 
         // build the cagra index
-        let build_params = IndexParams::new().unwrap();
         let index =
             Index::build(&res, &build_params, &dataset).expect("failed to create cagra index");
 
@@ -159,4 +157,18 @@ mod tests {
         assert_eq!(neighbors_host[[2, 0]], 2);
         assert_eq!(neighbors_host[[3, 0]], 3);
     }
+
+    #[test]
+    fn test_cagra_index() {
+        let build_params = IndexParams::new().unwrap();
+        test_cagra(build_params);
+    }
+
+    #[test]
+    fn test_cagra_compression() {
+        use crate::cagra::CompressionParams;
+        let build_params = IndexParams::new().unwrap()
+            .set_compression(CompressionParams::new().unwrap());
+        test_cagra(build_params);
+    }
 }