From d1c0f1b31c00d2902b79b472772089c7fb50cdce Mon Sep 17 00:00:00 2001 From: Tamas Bela Feher Date: Wed, 9 Nov 2022 23:11:27 +0100 Subject: [PATCH] Refactored IVF-PQ Python API as a thin wrapper around the cpp objects --- .../pylibraft/pylibraft/neighbors/__init__.py | 2 +- .../pylibraft/neighbors/c_ivf_pq.pxd | 18 +- .../pylibraft/pylibraft/neighbors/ivf_pq.pyx | 493 +++++++++++------- .../pylibraft/pylibraft/test/test_ivf_pq.py | 116 +++-- 4 files changed, 386 insertions(+), 243 deletions(-) diff --git a/python/pylibraft/pylibraft/neighbors/__init__.py b/python/pylibraft/pylibraft/neighbors/__init__.py index 3774d48a2a..aa9f8304fa 100644 --- a/python/pylibraft/pylibraft/neighbors/__init__.py +++ b/python/pylibraft/pylibraft/neighbors/__init__.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from .ivf_pq import IvfPq +# from .ivf_pq import IvfPq diff --git a/python/pylibraft/pylibraft/neighbors/c_ivf_pq.pxd b/python/pylibraft/pylibraft/neighbors/c_ivf_pq.pxd index 91e1eae836..dfc416b054 100644 --- a/python/pylibraft/pylibraft/neighbors/c_ivf_pq.pxd +++ b/python/pylibraft/pylibraft/neighbors/c_ivf_pq.pxd @@ -78,6 +78,18 @@ cdef extern from "raft/neighbors/ivf_pq_types.hpp" \ uint32_t pq_dim, uint32_t n_nonempty_lists) + IdxT size() + uint32_t dim() + uint32_t pq_dim() + uint32_t pq_len() + uint32_t pq_bits() + DistanceType metric() + uint32_t n_lists() + uint32_t rot_dim() + + + + cpdef cppclass search_params(ann_search_params): uint32_t n_probes cudaDataType_t lut_dtype @@ -90,19 +102,19 @@ cdef extern from "raft/neighbors/specializations/ivf_pq_specialization.hpp" \ const index_params& params, const float* dataset, uint64_t n_rows, - uint32_t dim) # except + + uint32_t dim) #except + cdef index[uint64_t] build(const handle_t& handle, const index_params& params, const int8_t* dataset, uint64_t n_rows, - uint32_t dim) # except + + uint32_t dim) #except + cdef index[uint64_t] build(const handle_t& handle, const index_params& params, const uint8_t* dataset, uint64_t n_rows, - uint32_t dim) # except + + uint32_t dim) #except + cdef index[uint64_t] extend(const handle_t& handle, const index[uint64_t]& orig_index, diff --git a/python/pylibraft/pylibraft/neighbors/ivf_pq.pyx b/python/pylibraft/pylibraft/neighbors/ivf_pq.pyx index 5ca75da794..6b2e943929 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_pq.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_pq.pyx @@ -24,9 +24,10 @@ import pylibraft.common.handle from libc.stdint cimport uintptr_t from libc.stdint cimport uint32_t, uint8_t, int8_t, int64_t, uint64_t from cython.operator cimport dereference as deref - from libcpp cimport bool, nullptr + from pylibraft.distance.distance_type cimport DistanceType +from pylibraft.common import Handle from pylibraft.common.handle cimport handle_t from rmm._lib.memory_resource cimport device_memory_resource @@ -68,21 +69,22 @@ def _check_input_array(cai, exp_dt, exp_rows=None, exp_cols=None): if exp_rows is not None and cai["shape"][0] != exp_rows: raise ValueError("Incorrect number of rows, expected {} , got {}" \ .format(exp_rows, cai["shape"][0])) - -class IvfPq: - """ - Nearest neighbors search using IVF-PQ method. - """ - # Class variables to provide easier access for data type parameters for the search function. - CUDA_R_32F = c_ivf_pq.cudaDataType_t.CUDA_R_32F - CUDA_R_16F = c_ivf_pq.cudaDataType_t.CUDA_R_16F - CUDA_R_8U = c_ivf_pq.cudaDataType_t.CUDA_R_8U +# Variables to provide easier access for parameters +PER_SUBSPACE = c_ivf_pq.codebook_gen.PER_SUBSPACE +PER_CLUSTER = c_ivf_pq.codebook_gen.PER_CLUSTER + +CUDA_R_32F = c_ivf_pq.cudaDataType_t.CUDA_R_32F +CUDA_R_16F = c_ivf_pq.cudaDataType_t.CUDA_R_16F +CUDA_R_8U = c_ivf_pq.cudaDataType_t.CUDA_R_8U + + +cdef class IndexParams: + cdef c_ivf_pq.index_params params def __init__(self, *, - handle=None, - n_lists = 1024, + n_lists=1024, metric="l2_expanded", kmeans_n_iters=20, kmeans_trainset_fraction=0.5, @@ -92,7 +94,7 @@ class IvfPq: force_random_rotation=False, add_data_on_build=True): """" - Approximate nearest neighbor search using IVF-PQ method. + Parameters to build index for IVF-PQ nearest neighbor search Parameters ---------- @@ -134,204 +136,329 @@ class IvfPq: After training the coarse and fine quantizers, we will populate the index with the dataset if add_data_on_build == True, otherwise the index is left empty, and the extend method can be used to add new vectors to the index. - - Examples - -------- - .. code-block:: python + """ + self.params.n_lists = n_lists + self.params.metric = _get_metric(metric) + self.params.metric_arg = 0 + self.params.kmeans_n_iters = kmeans_n_iters + self.params.kmeans_trainset_fraction = kmeans_trainset_fraction + self.params.pq_bits = pq_bits + self.params.pq_dim = pq_dim + if codebook_kind == "per_subspace": + self.params.codebook_kind = c_ivf_pq.codebook_gen.PER_SUBSPACE + elif codebook_kind == "per_cluster": + self.params.codebook_kind = c_ivf_pq.codebook_gen.PER_SUBSPACE + else: + raise ValueError("Incorrect codebook kind %s" % codebook_kind) + self.params.force_random_rotation = force_random_rotation + self.params.add_data_on_build = add_data_on_build + + @property + def n_lists(self): + return self.params.n_lists + + @property + def metric(self): + return self.params.metric + + @property + def kmeans_n_iters(self): + return self.params.kmeans_n_iters + + @property + def kmeans_trainset_fraction(self): + return self.params.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 codebook_kind(self): + return self.params.codebook_kind + + @property + def force_random_rotation(self): + return self.params.force_random_rotation + + @property + def add_data_on_build(self): + return self.params.add_data_on_build + + +cdef class Index: + # We store a pointer to the index because it dose not have a trivial constructor. + cdef c_ivf_pq.index[uint64_t] * index + cdef readonly bool trained + + def __cinit__(self, handle=None): + self.trained = False + self.index = NULL + if handle is None: + handle = Handle() + cdef handle_t* handle_ = handle.getHandle() + + # We create a placeholder object. The actual parameter values do not matter, it will be + # replaced with a built index object later. + self.index = new c_ivf_pq.index[uint64_t](deref(handle_), + _get_metric("l2_expanded"), + c_ivf_pq.codebook_gen.PER_SUBSPACE, + 1, + 4, + 8, + 0, + 0) + + def __dealloc__(self): + if self.index is not NULL: + del self.index + + @property + def dim(self): + return self.index[0].dim() + + @property + def size(self): + return self.index[0].size() + + @property + def pq_dim(self): + return self.index[0].pq_dim() + + @property + def pq_len(self): + return self.index[0].pq_len() + + @property + def pq_bits(self): + return self.index[0].pq_bits() + + @property + def metric(self): + return self.index[0].metric() - import cupy as cp + @property + def n_lists(self): + return self.index[0].n_lists() - from pylibraft.neighbors import IvfPq + @property + def rot_dim(self): + return self.index[0].rot_dim() - n_samples = 5000 - n_features = 50 - n_queries = 1000 - dataset = cp.random.random_sample((n_samples, n_features), +def build(IndexParams index_params, dataset, handle=None): + """ + Builds an IVF-PQ index that can be later used for nearest neighbor search. + + Parameters + ---------- + index_params : IndexParams object + dataset : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [float, int8, uint8] + + Examples + -------- + + .. code-block:: python + + import cupy as cp + + from pylibraft.common import Handle + from pylibraft.neighbors import ivf_pq + + n_samples = 5000 + n_features = 50 + n_queries = 1000 + + dataset = cp.random.random_sample((n_samples, n_features), dtype=cp.float32) - queries = cp.random.random_sample((n_samples, n_features), + queries = cp.random.random_sample((n_samples, n_features), dtype=cp.float32) - out_idx = cp.empty((n_samples, n_samples), dtype=cp.uint64) - out_dist = cp.empty((n_samples, n_samples), dtype=cp.float32) + out_idx = cp.empty((n_samples, n_samples), dtype=cp.uint64) + out_dist = cp.empty((n_samples, n_samples), dtype=cp.float32) - nn = IvfPQ() - nn.build(dataset) - [out_idx, out_dist] = nn.search(queries) - """ - self.handle = pylibraft.common.handle.Handle() if handle is None \ - else handle - - self._n_lists = n_lists - self._metric = metric - self._kmeans_n_iters = kmeans_n_iters - self._kmeans_trainset_fraction = kmeans_trainset_fraction - self._pq_bits = pq_bits - self._pq_dim = pq_dim - self._codebook_kind = codebook_kind - self._force_random_rotation = force_random_rotation - self._add_data_on_build = add_data_on_build - - self._index = None - - def __del__(self): - self._dealloc() - - def _dealloc(self): - # deallocate the index - cdef c_ivf_pq.index[uint64_t] *idx - if self._index is not None: - idx = self._index - del idx - - def build(self, dataset): - """ - Builds an IVF-PQ index that can be later used for nearest neighbor search. + handle = Handle() + build_params = ivf_pq.IndexParams() + index = build(index_params, dataset, handle) + search_params = ivf_pq.SearchParams() + [out_idx, out_dist] = ivf_pq.search(search_params, queries, handle) - Parameters - ---------- - dataset : CUDA array interface compliant matrix shape (n_samples, dim) - Supported dtype [float, int8, uint8] + # pylibraft functions are often asynchronous so the + # handle needs to be explicitly synchronized + handle.sync() - """ - # TODO(tfeher): ensure that this works with managed memory as well - dataset_cai = dataset.__cuda_array_interface__ - dataset_dt = np.dtype(dataset_cai["typestr"]) - _check_input_array(dataset_cai, [np.dtype('float32'), np.dtype('byte'), np.dtype('ubyte')]) + """ + dataset_cai = dataset.__cuda_array_interface__ + dataset_dt = np.dtype(dataset_cai["typestr"]) + _check_input_array(dataset_cai, [np.dtype('float32'), np.dtype('byte'), np.dtype('ubyte')]) + cdef uintptr_t dataset_ptr = dataset_cai["data"][0] - cdef c_ivf_pq.index_params params - params.n_lists = self._n_lists - params.metric = _get_metric(self._metric) - params.metric_arg = 0 - params.kmeans_n_iters = self._kmeans_n_iters - params.kmeans_trainset_fraction = self._kmeans_trainset_fraction - params.pq_bits = self._pq_bits - params.pq_dim = self._pq_dim - if self._codebook_kind == "per_subspace": - params.codebook_kind = c_ivf_pq.codebook_gen.PER_SUBSPACE - elif self._codebook_kind == "per_cluster": - params.codebook_kind = c_ivf_pq.codebook_gen.PER_SUBSPACE - else: - raise ValueError("Incorrect codebook kind %s" % self._codebook_kind) - params.force_random_rotation = self._force_random_rotation - params.add_data_on_build = self._add_data_on_build - - - # cdef index[uint64_t] *index_ptr - cdef uint64_t n_rows = dataset_cai["shape"][0] # make it uint32_t - self._dim = dataset_cai["shape"][1] - cdef handle_t* handle_ = self.handle.getHandle() - cdef uintptr_t dataset_ptr = dataset_cai["data"][0] + cdef uint64_t n_rows = dataset_cai["shape"][0] + cdef uint32_t dim = dataset_cai["shape"][1] + + if handle is None: + handle = Handle() + cdef handle_t* handle_ = handle.getHandle() - self._dealloc() - - cdef c_ivf_pq.index[uint64_t] *idx = new c_ivf_pq.index[uint64_t](deref(handle_), - _get_metric(self._metric), - c_ivf_pq.codebook_gen.PER_SUBSPACE, - self._n_lists, - self._dim, - self._pq_bits, - self._pq_dim, - 0) + idx = Index() - if dataset_dt == np.float32: - idx[0] = c_ivf_pq.build(deref(handle_), - params, + if dataset_dt == np.float32: + idx.index[0] = c_ivf_pq.build(deref(handle_), + index_params.params, dataset_ptr, n_rows, - self._dim) - elif dataset_dt == np.byte: - idx[0] = c_ivf_pq.build(deref(handle_), - params, + dim) + idx.trained = True + elif dataset_dt == np.byte: + idx.index[0] = c_ivf_pq.build(deref(handle_), + index_params.params, dataset_ptr, n_rows, - self._dim) - elif dataset_dt == np.ubyte: - idx[0] = c_ivf_pq.build(deref(handle_), - params, + dim) + idx.trained = True + elif dataset_dt == np.ubyte: + idx.index[0] = c_ivf_pq.build(deref(handle_), + index_params.params, dataset_ptr, n_rows, - self._dim) - else: - raise TypeError("dtype %s not supported" % dataset_dt) - - self._index = idx + dim) + idx.trained = True + else: + raise TypeError("dtype %s not supported" % dataset_dt) + + handle.sync() + return idx - self.handle.sync() - def extend(self, new_vectors, new_indices): - """ - Extend an existing index with new vectors. +def extend(Index index, new_vectors, new_indices, handle=None): + """ + Extend an existing index with new vectors. - Parameters - ---------- - new_vectors : CUDA array interface compliant matrix shape (n_samples, dim) - Supported dtype [float, int8, uint8] - new_indices : CUDA array interface compliant matrix shape (n_samples, dim) - Supported dtype [uint64] - """ - if self._index is None: - raise ValueError("Index need to be built before calling extend.") + Parameters + ---------- + index : ivf_pq.Index + Trained ivf_pq object. + new_vectors : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [float, int8, uint8] + new_indices : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [uint64] + handle: raft Handle + + """ + if not index.trained: + raise ValueError("Index need to be built before calling extend.") - vecs_cai = new_vectors.__cuda_array_interface__ - vecs_dt = np.dtype(vecs_cai["typestr"]) - cdef uint32_t n_rows = vecs_cai["shape"][0] - cdef uint32_t dim = vecs_cai["shape"][1] + if handle is None: + handle = Handle() + cdef handle_t* handle_ = handle.getHandle() - _check_input_array(vecs_cai, [np.dtype('float32'), np.dtype('byte'), np.dtype('ubyte')], - exp_cols=self._dim) + vecs_cai = new_vectors.__cuda_array_interface__ + vecs_dt = np.dtype(vecs_cai["typestr"]) + cdef uint64_t n_rows = vecs_cai["shape"][0] + cdef uint32_t dim = vecs_cai["shape"][1] - idx_cai = new_indices.__cuda_array_interface__ - _check_input_array(idx_cai, [np.dtype('uint64')], exp_rows=n_rows) - if len(idx_cai["shape"])!=1: - raise ValueError("Indices array is expected to be 1D") + _check_input_array(vecs_cai, [np.dtype('float32'), np.dtype('byte'), np.dtype('ubyte')], + exp_cols=index.dim) + idx_cai = new_indices.__cuda_array_interface__ + _check_input_array(idx_cai, [np.dtype('uint64')], exp_rows=n_rows) + if len(idx_cai["shape"])!=1: + raise ValueError("Indices array is expected to be 1D") - cdef handle_t* handle_ = self.handle.getHandle() - cdef c_ivf_pq.index[uint64_t] *idx = self._index - cdef uintptr_t vecs_ptr = vecs_cai["data"][0] - cdef uintptr_t idx_ptr = idx_cai["data"][0] - if vecs_dt == np.float32: - idx[0] = c_ivf_pq.extend(deref(handle_), - deref(idx), + cdef uintptr_t vecs_ptr = vecs_cai["data"][0] + cdef uintptr_t idx_ptr = idx_cai["data"][0] + + if vecs_dt == np.float32: + index.index[0] = c_ivf_pq.extend(deref(handle_), + deref(index.index), vecs_ptr, idx_ptr, n_rows) - elif vecs_dt == np.int8: - idx[0] = c_ivf_pq.extend(deref(handle_), - deref(idx), + elif vecs_dt == np.int8: + index.index[0] = c_ivf_pq.extend(deref(handle_), + deref(index.index), vecs_ptr, idx_ptr, n_rows) - elif vecs_dt == np.uint8: - idx[0] = c_ivf_pq.extend(deref(handle_), - deref(idx), + elif vecs_dt == np.uint8: + index.index[0] = c_ivf_pq.extend(deref(handle_), + deref(index.index), vecs_ptr, idx_ptr, n_rows) - else: - raise TypeError("query dtype %s not supported" % vecs_dt) + else: + raise TypeError("query dtype %s not supported" % vecs_dt) + + handle.sync() + + return index + - self.handle.sync() +cdef class SearchParams: + cdef c_ivf_pq.search_params params + def __init__(self, *, n_probes=20, + lut_dtype=CUDA_R_32F, + internal_distance_dtype=CUDA_R_32F): + """ + IVF-PQ search parameters + Parameters + ---------- + n_probes: int, default = 1024 + The number of course clusters to select for the fine search. + lut_dtype: default = ivf_pq.CUDA_R_32F (float) + Data type of look up table to be created dynamically at search time. The use of + low-precision types reduces the amount of shared memory required at search time, so + fast shared memory kernels can be used even for datasets with large dimansionality. + Note that the recall is slightly degraded when low-precision type is selected. + Possible values [CUDA_R_32F, CUDA_R_16F, CUDA_R_8U] + internal_distance_dtype: default = ivf_q.CUDA_R_32F (float) + Storage data type for distance/similarity computation. + Possible values [CUDA_R_32F, CUDA_R_16F] + + """ - def search(self, - queries, - k, - neighbors, - distances, - n_probes=20, - lut_dtype=CUDA_R_32F, - internal_distance_dtype=CUDA_R_32F, - preferred_thread_block_size=0): + self.params.n_probes = n_probes + self.params.lut_dtype = lut_dtype + self.params.internal_distance_dtype = internal_distance_dtype + # self.params.shmem_carveout = self.shmem_carveout # TODO(tfeher): enable if #926 adds this + + @property + def n_probes(self): + return self.params.n_probes + + @property + def lut_dtype(self): + return self.params.lut_dtype + + @property + def internal_distance_dtype(self): + return self.params.internal_distance_dtype + +def search(SearchParams search_params, + Index index, + queries, + k, + neighbors, + distances, + handle=None + ): """ Find the k nearest neighbors for each query. Parameters ---------- + search_params : SearchParams + index : Index + Trained IVF-PQ index. queries : CUDA array interface compliant matrix shape (n_samples, dim) Supported dtype [float, int8, uint8] k : int @@ -342,33 +469,26 @@ class IvfPq: distances : CUDA array interface compliant matrix shape (n_queries, k) If this parameter is specified, then the distances to the neighbors will be returned here. Otherwise a new array is created. - n_probes: int, default = 1024 - The number of course clusters to select for the fine search. - lut_dtype: default = IvfPq.CUDA_R_32F (float) - Data type of look up table to be created dynamically at search time. The use of - low-precision types reduces the amount of shared memory required at search time, so - fast shared memory kernels can be used even for datasets with large dimansionality. - Note that the recall is slightly degraded when low-precision type is selected. - Possible values [CUDA_R_32F, CUDA_R_16F, CUDA_R_8U] - internal_distance_dtype: default = IvfPq.CUDA_R_32F (float) - Storage data type for distance/similarity computation. - Possible values [CUDA_R_32F, CUDA_R_16F] Returns ------- - A pair of [neighbors, distances] arrays as defined above. + A pair of (neighbors, distances) arrays as defined above. """ - if self._index is None: + if not index.trained: raise ValueError("Index need to be built before calling search.") + if handle is None: + handle = Handle() + cdef handle_t* handle_ = handle.getHandle() + queries_cai = queries.__cuda_array_interface__ queries_dt = np.dtype(queries_cai["typestr"]) cdef uint32_t n_queries = queries_cai["shape"][0] _check_input_array(queries_cai, [np.dtype('float32'), np.dtype('byte'), np.dtype('ubyte')], - exp_cols=self._dim) + exp_cols=index.dim) neighbors_cai = neighbors.__cuda_array_interface__ _check_input_array(neighbors_cai, [np.dtype('uint64')], exp_rows=n_queries, exp_cols=k) @@ -376,13 +496,8 @@ class IvfPq: distances_cai = distances.__cuda_array_interface__ _check_input_array(distances_cai, [np.dtype('float32')], exp_rows=n_queries, exp_cols=k) - cdef c_ivf_pq.search_params params - params.n_probes = n_probes - params.lut_dtype = lut_dtype - params.internal_distance_dtype = internal_distance_dtype + cdef c_ivf_pq.search_params params = search_params.params - cdef handle_t* handle_ = self.handle.getHandle() - cdef c_ivf_pq.index[uint64_t] *idx = self._index cdef uintptr_t queries_ptr = queries_cai["data"][0] cdef uintptr_t neighbors_ptr = neighbors_cai["data"][0] cdef uintptr_t distances_ptr = distances_cai["data"][0] @@ -391,7 +506,7 @@ class IvfPq: if queries_dt == np.float32: c_ivf_pq.search(deref(handle_), params, - deref(idx), + deref(index.index), queries_ptr, n_queries, k, @@ -401,7 +516,7 @@ class IvfPq: elif queries_dt == np.byte: c_ivf_pq.search(deref(handle_), params, - deref(idx), + deref(index.index), queries_ptr, n_queries, k, @@ -411,7 +526,7 @@ class IvfPq: elif queries_dt == np.ubyte: c_ivf_pq.search(deref(handle_), params, - deref(idx), + deref(index.index), queries_ptr, n_queries, k, @@ -421,5 +536,7 @@ class IvfPq: else: raise ValueError("query dtype %s not supported" % queries_dt) - self.handle.sync() + handle.sync() + + return (neighbors, distances) diff --git a/python/pylibraft/pylibraft/test/test_ivf_pq.py b/python/pylibraft/pylibraft/test/test_ivf_pq.py index c94e2e44c3..dc83a9ad83 100644 --- a/python/pylibraft/pylibraft/test/test_ivf_pq.py +++ b/python/pylibraft/pylibraft/test/test_ivf_pq.py @@ -19,7 +19,7 @@ from sklearn.metrics import pairwise_distances from sklearn.preprocessing import normalize -from pylibraft.neighbors import IvfPq +from pylibraft.neighbors import ivf_pq from pylibraft.testing.utils import TestDeviceBuffer @@ -85,8 +85,8 @@ def run_ivf_pq_build_search_test( codebook_kind="per_cluster", add_data_on_build="True", n_probes=100, - lut_dtype=IvfPq.CUDA_R_32F, - internal_distance_dtype=IvfPq.CUDA_R_32F, + lut_dtype=ivf_pq.CUDA_R_32F, + internal_distance_dtype=ivf_pq.CUDA_R_32F, force_random_rotation=False, kmeans_trainset_fraction=1, kmeans_n_iters=20, @@ -97,7 +97,7 @@ def run_ivf_pq_build_search_test( dataset = normalize(dataset, norm="l2", axis=1) dataset_device = TestDeviceBuffer(dataset, order="C") - nn = IvfPq( + build_params = ivf_pq.IndexParams( n_lists=n_lists, metric=metric, kmeans_n_iters=kmeans_n_iters, @@ -109,9 +109,14 @@ def run_ivf_pq_build_search_test( add_data_on_build=add_data_on_build, ) - nn.build(dataset_device) + index = ivf_pq.build(build_params, dataset_device) - assert nn._index is not None + assert index.trained + if pq_dim != 0: + assert index.pq_dim == build_params.pq_dim + assert index.pq_bits == build_params.pq_bits + assert index.metric == build_params.metric + assert index.n_lists == build_params.n_lists if not add_data_on_build: dataset_1_device = TestDeviceBuffer(dataset[: n_rows // 2, :], order="C") @@ -120,8 +125,8 @@ def run_ivf_pq_build_search_test( indices_1_device = TestDeviceBuffer(indices_1, order="C") indices_2 = np.arange(n_rows // 2, n_rows, dtype=np.uint64) indices_2_device = TestDeviceBuffer(indices_2, order="C") - nn.extend(dataset_1_device, indices_1_device) - nn.extend(dataset_2_device, indices_2_device) + index = ivf_pq.extend(index, dataset_1_device, indices_1_device) + index = ivf_pq.extend(index, dataset_2_device, indices_2_device) queries = generate_data((n_queries, n_cols), dtype) out_idx = np.zeros((n_queries, k), dtype=np.uint64) @@ -131,14 +136,19 @@ def run_ivf_pq_build_search_test( out_idx_device = TestDeviceBuffer(out_idx, order="C") out_dist_device = TestDeviceBuffer(out_dist, order="C") - nn.search( + search_params = ivf_pq.SearchParams( + n_probes=n_probes, + lut_dtype=lut_dtype, + internal_distance_dtype=internal_distance_dtype, + ) + + ivf_pq.search( + search_params, + index, queries_device, k, out_idx_device, out_dist_device, - n_probes=n_probes, - lut_dtype=lut_dtype, - internal_distance_dtype=internal_distance_dtype, ) if not compare: @@ -178,26 +188,26 @@ def test_ivf_pq_dtypes(n_rows, n_cols, n_queries, n_lists, dtype): ) -@pytest.mark.parametrize( - "params", - [ - {"n_rows": 1, "n_cols": 10, "n_queries": 10, "k": 1, "n_lists": 10}, - {"n_rows": 10, "n_cols": 1, "n_queries": 10, "k": 10, "n_lists": 10}, - {"n_rows": 999, "n_cols": 42, "n_queries": 4953, "k": 137, "n_lists": 53}, - ], -) -def test_ivf_pq_n(params): - # We do not test recall, just confirm that we can handle edge cases for certain parameters - run_ivf_pq_build_search_test( - n_rows=params["n_rows"], - n_cols=params["n_cols"], - n_queries=params["n_queries"], - k=params["k"], - n_lists=params["n_lists"], - metric="l2_expanded", - dtype=np.float32, - compare=False, - ) +# @pytest.mark.parametrize( +# "params", +# [ +# {"n_rows": 1, "n_cols": 10, "n_queries": 10, "k": 1, "n_lists": 10}, +# {"n_rows": 10, "n_cols": 1, "n_queries": 10, "k": 10, "n_lists": 10}, +# {"n_rows": 999, "n_cols": 42, "n_queries": 4953, "k": 137, "n_lists": 53}, +# ], +# ) +# def test_ivf_pq_n(params): +# # We do not test recall, just confirm that we can handle edge cases for certain parameters +# run_ivf_pq_build_search_test( +# n_rows=params["n_rows"], +# n_cols=params["n_cols"], +# n_queries=params["n_queries"], +# k=params["k"], +# n_lists=params["n_lists"], +# metric="l2_expanded", +# dtype=np.float32, +# compare=False, +# ) @pytest.mark.parametrize("metric", ["l2_expanded", "inner_product"]) @@ -256,10 +266,10 @@ def test_ivf_pq_params(params): @pytest.mark.parametrize( "params", [ - {"k": 10, "n_probes": 100, "lut": IvfPq.CUDA_R_16F, "idd": IvfPq.CUDA_R_32F}, - {"k": 10, "n_probes": 99, "lut": IvfPq.CUDA_R_8U, "idd": IvfPq.CUDA_R_32F}, - {"k": 10, "n_probes": 100, "lut": IvfPq.CUDA_R_32F, "idd": IvfPq.CUDA_R_16F}, - {"k": 129, "n_probes": 100, "lut": IvfPq.CUDA_R_32F, "idd": IvfPq.CUDA_R_32F}, + {"k": 10, "n_probes": 100, "lut": ivf_pq.CUDA_R_16F, "idd": ivf_pq.CUDA_R_32F}, + {"k": 10, "n_probes": 99, "lut": ivf_pq.CUDA_R_8U, "idd": ivf_pq.CUDA_R_32F}, + {"k": 10, "n_probes": 100, "lut": ivf_pq.CUDA_R_32F, "idd": ivf_pq.CUDA_R_16F}, + {"k": 129, "n_probes": 100, "lut": ivf_pq.CUDA_R_32F, "idd": ivf_pq.CUDA_R_32F}, ], ) def test_ivf_pq_search_params(params): @@ -310,7 +320,7 @@ def test_build_assertions(): dataset = generate_data((n_rows, n_cols), np.float32) dataset_device = TestDeviceBuffer(dataset, order="C") - nn = IvfPq( + index_params = ivf_pq.IndexParams( n_lists=50, metric="l2_expanded", kmeans_n_iters=20, @@ -318,6 +328,8 @@ def test_build_assertions(): add_data_on_build=False, ) + index = ivf_pq.Index() + queries = generate_data((n_queries, n_cols), np.float32) out_idx = np.zeros((n_queries, k), dtype=np.uint64) out_dist = np.zeros((n_queries, k), dtype=np.float32) @@ -326,23 +338,27 @@ def test_build_assertions(): out_idx_device = TestDeviceBuffer(out_idx, order="C") out_dist_device = TestDeviceBuffer(out_dist, order="C") + search_params = ivf_pq.SearchParams(n_probes=50) + with pytest.raises(ValueError): # Index must be built before search - nn.search(queries_device, k, out_idx_device, out_dist_device, n_probes=50) + ivf_pq.search( + search_params, index, queries_device, k, out_idx_device, out_dist_device + ) - nn.build(dataset_device) - assert nn._index is not None + index = ivf_pq.build(index_params, dataset_device) + assert index.trained indices = np.arange(n_rows + 1, dtype=np.uint64) indices_device = TestDeviceBuffer(indices, order="C") with pytest.raises(ValueError): # Dataset dimension mismatch - nn.extend(queries_device, indices_device) + ivf_pq.extend(index, queries_device, indices_device) with pytest.raises(ValueError): # indices dimension mismatch - nn.extend(dataset_device, indices_device) + ivf_pq.extend(index, dataset_device, indices_device) @pytest.mark.parametrize( @@ -394,17 +410,15 @@ def test_search_inputs(params): ) out_dist_device = TestDeviceBuffer(out_dist, order=dist_order) - nn = IvfPq(n_lists=50, metric="l2_expanded", add_data_on_build=True) + index_params = ivf_pq.IndexParams( + n_lists=50, metric="l2_expanded", add_data_on_build=True + ) dataset = generate_data((n_rows, n_cols), dtype) dataset_device = TestDeviceBuffer(dataset, order="C") - nn.build(dataset_device) - assert nn._index is not None + index = ivf_pq.build(index_params, dataset_device) + assert index.trained with pytest.raises(Exception): - nn.search(queries_device, k, out_idx_device, out_dist_device, n_probes=50) - - -def test_new_api(): - params = IvfPq.index_params - assert params.n_litst > 0 + search_params = ivf_pq.SearchParams(n_probes=50) + search(search_params, index, queries_device, k, out_idx_device, out_dist_device)