From cf405695eccc10f1b1191b0a134d2f06e5afdf50 Mon Sep 17 00:00:00 2001 From: Tamas Bela Feher Date: Tue, 1 Nov 2022 02:32:35 +0100 Subject: [PATCH] Work in progress, build defined for float inputs --- cpp/CMakeLists.txt | 124 +++---- .../specializations/ivf_pq_specialization.hpp | 33 +- .../nn/specializations/detail/ivfpq_build.cu | 46 +++ .../nn/specializations/detail/ivfpq_search.cu | 39 ++ .../pylibraft/neighbors/CMakeLists.txt | 2 +- .../pylibraft/pylibraft/neighbors/__init__.py | 2 +- .../pylibraft/pylibraft/neighbors/ivf_pq.pyx | 339 ++++++++++++++++++ .../pylibraft/neighbors/neighbors.pyx | 89 ----- 8 files changed, 517 insertions(+), 157 deletions(-) create mode 100644 cpp/src/nn/specializations/detail/ivfpq_build.cu create mode 100644 cpp/src/nn/specializations/detail/ivfpq_search.cu create mode 100644 python/pylibraft/pylibraft/neighbors/ivf_pq.pyx delete mode 100644 python/pylibraft/pylibraft/neighbors/neighbors.pyx diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 3a60d8a28f..f2928ac2c2 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -245,57 +245,57 @@ if(RAFT_COMPILE_DIST_LIBRARY) add_library(raft_distance_lib src/distance/pairwise_distance.cu src/distance/fused_l2_min_arg.cu -# src/distance/specializations/detail/canberra.cu -# src/distance/specializations/detail/chebyshev.cu -# src/distance/specializations/detail/correlation.cu -# src/distance/specializations/detail/cosine.cu -# src/distance/specializations/detail/cosine.cu -# src/distance/specializations/detail/hamming_unexpanded.cu -# src/distance/specializations/detail/hellinger_expanded.cu -# src/distance/specializations/detail/jensen_shannon_float_float_float_int.cu -# src/distance/specializations/detail/jensen_shannon_float_float_float_uint32.cu -# src/distance/specializations/detail/jensen_shannon_double_double_double_int.cu -# src/distance/specializations/detail/kernels/gram_matrix_base_double.cu -# src/distance/specializations/detail/kernels/gram_matrix_base_float.cu -# src/distance/specializations/detail/kernels/polynomial_kernel_double_int.cu -# src/distance/specializations/detail/kernels/polynomial_kernel_float_int.cu -# # These are somehow missing a kernel definition which is causing a compile error. -# # src/distance/specializations/detail/kernels/rbf_kernel_double.cu -# # src/distance/specializations/detail/kernels/rbf_kernel_float.cu -# src/distance/specializations/detail/kernels/tanh_kernel_double.cu -# src/distance/specializations/detail/kernels/tanh_kernel_float.cu -# src/distance/specializations/detail/kl_divergence_float_float_float_int.cu -# src/distance/specializations/detail/kl_divergence_float_float_float_uint32.cu -# src/distance/specializations/detail/kl_divergence_double_double_double_int.cu -# src/distance/specializations/detail/l1_float_float_float_int.cu -# src/distance/specializations/detail/l1_float_float_float_uint32.cu -# src/distance/specializations/detail/l1_double_double_double_int.cu + src/distance/specializations/detail/canberra.cu + src/distance/specializations/detail/chebyshev.cu + src/distance/specializations/detail/correlation.cu + src/distance/specializations/detail/cosine.cu + src/distance/specializations/detail/cosine.cu + src/distance/specializations/detail/hamming_unexpanded.cu + src/distance/specializations/detail/hellinger_expanded.cu + src/distance/specializations/detail/jensen_shannon_float_float_float_int.cu + src/distance/specializations/detail/jensen_shannon_float_float_float_uint32.cu + src/distance/specializations/detail/jensen_shannon_double_double_double_int.cu + src/distance/specializations/detail/kernels/gram_matrix_base_double.cu + src/distance/specializations/detail/kernels/gram_matrix_base_float.cu + src/distance/specializations/detail/kernels/polynomial_kernel_double_int.cu + src/distance/specializations/detail/kernels/polynomial_kernel_float_int.cu +# These are somehow missing a kernel definition which is causing a compile error. +# src/distance/specializations/detail/kernels/rbf_kernel_double.cu +# src/distance/specializations/detail/kernels/rbf_kernel_float.cu + src/distance/specializations/detail/kernels/tanh_kernel_double.cu + src/distance/specializations/detail/kernels/tanh_kernel_float.cu + src/distance/specializations/detail/kl_divergence_float_float_float_int.cu + src/distance/specializations/detail/kl_divergence_float_float_float_uint32.cu + src/distance/specializations/detail/kl_divergence_double_double_double_int.cu + src/distance/specializations/detail/l1_float_float_float_int.cu + src/distance/specializations/detail/l1_float_float_float_uint32.cu + src/distance/specializations/detail/l1_double_double_double_int.cu src/distance/specializations/detail/l2_expanded_float_float_float_int.cu src/distance/specializations/detail/l2_expanded_float_float_float_uint32.cu src/distance/specializations/detail/l2_expanded_double_double_double_int.cu - # src/distance/specializations/detail/l2_sqrt_expanded_float_float_float_int.cu - # src/distance/specializations/detail/l2_sqrt_expanded_float_float_float_uint32.cu - # src/distance/specializations/detail/l2_sqrt_expanded_double_double_double_int.cu - # src/distance/specializations/detail/l2_sqrt_unexpanded_float_float_float_int.cu - # src/distance/specializations/detail/l2_sqrt_unexpanded_float_float_float_uint32.cu - # src/distance/specializations/detail/l2_sqrt_unexpanded_double_double_double_int.cu - # src/distance/specializations/detail/l2_unexpanded_double_double_double_int.cu - # src/distance/specializations/detail/l2_unexpanded_float_float_float_uint32.cu - # src/distance/specializations/detail/l2_unexpanded_float_float_float_int.cu - # src/distance/specializations/detail/lp_unexpanded_double_double_double_int.cu - # src/distance/specializations/detail/lp_unexpanded_float_float_float_uint32.cu - # src/distance/specializations/detail/lp_unexpanded_float_float_float_int.cu - # src/distance/specializations/detail/russel_rao_double_double_double_int.cu - # src/distance/specializations/detail/russel_rao_float_float_float_uint32.cu - # src/distance/specializations/detail/russel_rao_float_float_float_int.cu - # src/distance/specializations/fused_l2_nn_double_int.cu - # src/distance/specializations/fused_l2_nn_double_int64.cu - # src/distance/specializations/fused_l2_nn_float_int.cu - # src/distance/specializations/fused_l2_nn_float_int64.cu - # src/random/specializations/rmat_rectangular_generator_int_double.cu - # src/random/specializations/rmat_rectangular_generator_int64_double.cu - # src/random/specializations/rmat_rectangular_generator_int_float.cu - # src/random/specializations/rmat_rectangular_generator_int64_float.cu + src/distance/specializations/detail/l2_sqrt_expanded_float_float_float_int.cu + src/distance/specializations/detail/l2_sqrt_expanded_float_float_float_uint32.cu + src/distance/specializations/detail/l2_sqrt_expanded_double_double_double_int.cu + src/distance/specializations/detail/l2_sqrt_unexpanded_float_float_float_int.cu + src/distance/specializations/detail/l2_sqrt_unexpanded_float_float_float_uint32.cu + src/distance/specializations/detail/l2_sqrt_unexpanded_double_double_double_int.cu + src/distance/specializations/detail/l2_unexpanded_double_double_double_int.cu + src/distance/specializations/detail/l2_unexpanded_float_float_float_uint32.cu + src/distance/specializations/detail/l2_unexpanded_float_float_float_int.cu + src/distance/specializations/detail/lp_unexpanded_double_double_double_int.cu + src/distance/specializations/detail/lp_unexpanded_float_float_float_uint32.cu + src/distance/specializations/detail/lp_unexpanded_float_float_float_int.cu + src/distance/specializations/detail/russel_rao_double_double_double_int.cu + src/distance/specializations/detail/russel_rao_float_float_float_uint32.cu + src/distance/specializations/detail/russel_rao_float_float_float_int.cu + src/distance/specializations/fused_l2_nn_double_int.cu + src/distance/specializations/fused_l2_nn_double_int64.cu + src/distance/specializations/fused_l2_nn_float_int.cu + src/distance/specializations/fused_l2_nn_float_int64.cu + src/random/specializations/rmat_rectangular_generator_int_double.cu + src/random/specializations/rmat_rectangular_generator_int64_double.cu + src/random/specializations/rmat_rectangular_generator_int_float.cu + src/random/specializations/rmat_rectangular_generator_int64_float.cu ) set_target_properties( raft_distance_lib @@ -354,23 +354,25 @@ if(RAFT_COMPILE_NN_LIBRARY) src/nn/specializations/detail/ivfpq_compute_similarity_float_fast.cu src/nn/specializations/detail/ivfpq_compute_similarity_float_no_basediff.cu src/nn/specializations/detail/ivfpq_compute_similarity_float_no_smem_lut.cu - # src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_fast.cu - # src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_no_basediff.cu - # src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_no_smem_lut.cu - # src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_fast.cu - # src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_no_basediff.cu - # src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_no_smem_lut.cu - # src/nn/specializations/detail/ivfpq_compute_similarity_half_fast.cu - # src/nn/specializations/detail/ivfpq_compute_similarity_half_no_basediff.cu - # src/nn/specializations/detail/ivfpq_compute_similarity_half_no_smem_lut.cu - # src/nn/specializations/detail/ivfpq_search_float_int64_t.cu - # src/nn/specializations/detail/ivfpq_search_float_uint32_t.cu - # src/nn/specializations/detail/ivfpq_search_float_uint64_t.cu + src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_fast.cu + src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_no_basediff.cu + src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_no_smem_lut.cu + src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_fast.cu + src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_no_basediff.cu + src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_no_smem_lut.cu + src/nn/specializations/detail/ivfpq_compute_similarity_half_fast.cu + src/nn/specializations/detail/ivfpq_compute_similarity_half_no_basediff.cu + src/nn/specializations/detail/ivfpq_compute_similarity_half_no_smem_lut.cu + src/nn/specializations/detail/ivfpq_build.cu + src/nn/specializations/detail/ivfpq_search.cu + src/nn/specializations/detail/ivfpq_search_float_int64_t.cu + src/nn/specializations/detail/ivfpq_search_float_uint32_t.cu + src/nn/specializations/detail/ivfpq_search_float_uint64_t.cu src/nn/specializations/fused_l2_knn_long_float_true.cu src/nn/specializations/fused_l2_knn_long_float_false.cu src/nn/specializations/fused_l2_knn_int_float_true.cu src/nn/specializations/fused_l2_knn_int_float_false.cu - src/nn/specializations/knn.cu + # src/nn/specializations/knn.cu ) set_target_properties( raft_nn_lib diff --git a/cpp/include/raft/neighbors/specializations/ivf_pq_specialization.hpp b/cpp/include/raft/neighbors/specializations/ivf_pq_specialization.hpp index 3ce8237c89..b7e70636fb 100644 --- a/cpp/include/raft/neighbors/specializations/ivf_pq_specialization.hpp +++ b/cpp/include/raft/neighbors/specializations/ivf_pq_specialization.hpp @@ -20,7 +20,7 @@ namespace raft::neighbors ::ivf_pq { -#define RAFT_INST(T, IdxT) \ +#define RAFT_INST_SEARCH(T, IdxT) \ void search(const handle_t&, \ const search_params&, \ const index&, \ @@ -31,10 +31,33 @@ namespace raft::neighbors ::ivf_pq { float*, \ rmm::mr::device_memory_resource*); -RAFT_INST(float, int64_t); -RAFT_INST(float, uint32_t); -RAFT_INST(float, uint64_t); +RAFT_INST_SEARCH(float, int64_t); +RAFT_INST_SEARCH(float, uint32_t); +RAFT_INST_SEARCH(float, uint64_t); -#undef RAFT_INST +#undef RAFT_INST_SEARCH + +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + auto build(const handle_t& handle, \ + const index_params& params, \ + const T* dataset, \ + IdxT n_rows, \ + uint32_t dim) \ + ->index; \ + \ + auto extend(const handle_t& handle, \ + const index& orig_index, \ + const T* new_vectors, \ + const IdxT* new_indices, \ + IdxT n_rows) \ + ->index; + +RAFT_INST_BUILD_EXTEND(float, int64_t) +RAFT_INST_BUILD_EXTEND(int8_t, int64_t) +RAFT_INST_BUILD_EXTEND(uint8_t, int64_t) +// RAFT_INST_BUILD_EXTEND(float, uint32_t); +// RAFT_INST_BUILD_EXTEND(float, uint64_t); + +#undef RAFT_INST_BUILD_EXTEND } // namespace raft::neighbors::ivf_pq diff --git a/cpp/src/nn/specializations/detail/ivfpq_build.cu b/cpp/src/nn/specializations/detail/ivfpq_build.cu new file mode 100644 index 0000000000..474c42595b --- /dev/null +++ b/cpp/src/nn/specializations/detail/ivfpq_build.cu @@ -0,0 +1,46 @@ +/* + * Copyright (c) 2022, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +namespace raft::neighbors::ivf_pq { + +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + template auto build(const handle_t& handle, \ + const index_params& params, \ + const T* dataset, \ + IdxT n_rows, \ + uint32_t dim) \ + ->index; \ + \ + template auto extend(const handle_t& handle, \ + const index& orig_index, \ + const T* new_vectors, \ + const IdxT* new_indices, \ + IdxT n_rows) \ + ->index; + +RAFT_INST_BUILD_EXTEND(float, int64_t); +RAFT_INST_BUILD_EXTEND(int8_t, int64_t); +RAFT_INST_BUILD_EXTEND(uint8_t, int64_t); + +// RAFT_INST_BUILD_EXTEND(float, uint32_t); +// RAFT_INST_BUILD_EXTEND(float, uint64_t); + +#undef RAFT_INST_BUILD_EXTEND + +} // namespace raft::neighbors::ivf_pq diff --git a/cpp/src/nn/specializations/detail/ivfpq_search.cu b/cpp/src/nn/specializations/detail/ivfpq_search.cu new file mode 100644 index 0000000000..4679b35646 --- /dev/null +++ b/cpp/src/nn/specializations/detail/ivfpq_search.cu @@ -0,0 +1,39 @@ +/* + * Copyright (c) 2022, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +namespace raft::neighbors::ivf_pq { + +#define RAFT_SEARCH_INST(T, IdxT) \ + template void search(const handle_t&, \ + const search_params&, \ + const index&, \ + const T*, \ + uint32_t, \ + uint32_t, \ + IdxT*, \ + float*, \ + rmm::mr::device_memory_resource*); + +RAFT_SEARCH_INST(float, int64_t); +RAFT_SEARCH_INST(float, uint32_t); +RAFT_SEARCH_INST(float, uint64_t); + +#undef RAFT_INST_SEARCH + +} // namespace raft::neighbors::ivf_pq diff --git a/python/pylibraft/pylibraft/neighbors/CMakeLists.txt b/python/pylibraft/pylibraft/neighbors/CMakeLists.txt index 58fa83bd5a..0af059e084 100644 --- a/python/pylibraft/pylibraft/neighbors/CMakeLists.txt +++ b/python/pylibraft/pylibraft/neighbors/CMakeLists.txt @@ -13,7 +13,7 @@ # ============================================================================= # Set the list of Cython files to build -set(cython_sources neighbors.pyx) +set(cython_sources ivf_pq.pyx) set(linked_libraries raft::raft raft::nn) # Build all of the Cython targets diff --git a/python/pylibraft/pylibraft/neighbors/__init__.py b/python/pylibraft/pylibraft/neighbors/__init__.py index 527e644b8f..8aa907c7de 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 diff --git a/python/pylibraft/pylibraft/neighbors/ivf_pq.pyx b/python/pylibraft/pylibraft/neighbors/ivf_pq.pyx new file mode 100644 index 0000000000..537fb25ce3 --- /dev/null +++ b/python/pylibraft/pylibraft/neighbors/ivf_pq.pyx @@ -0,0 +1,339 @@ +# +# Copyright (c) 2022, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: profile=False +# distutils: language = c++ +# cython: embedsignature = True +# cython: language_level = 3 + +import numpy as np +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 cuda.ccudart cimport cudaDataType_t + +from libcpp cimport bool +from pylibraft.distance.distance_type cimport DistanceType +from pylibraft.common.handle cimport handle_t +from rmm._lib.memory_resource cimport device_memory_resource + +cdef extern from "raft/neighbors/ann_types.hpp" \ + namespace "raft::neighbors::ann": + + cdef cppclass ann_index "raft::neighbors::index": + pass + + cdef cppclass ann_index_params "raft::spatial::knn::index_params": + DistanceType metric + float metric_arg + bool add_data_on_build + + cdef cppclass ann_search_params "raft::spatial::knn::search_params": + pass + +cdef extern from "raft/neighbors/ivf_pq_types.hpp" \ + namespace "raft::neighbors::ivf_pq": + + ctypedef enum codebook_gen: + PER_SUBSPACE "raft::neighbors::ivf_pq::codebook_gen::PER_SUBSPACE", + PER_CLUSTER "raft::neighbors::ivf_pq::codebook_gen::PER_CLUSTER" + + + cdef cppclass index_params(ann_index_params): + uint32_t n_lists + uint32_t kmeans_n_iters + double kmeans_trainset_fraction + uint32_t pq_bits + uint32_t pq_dim + codebook_gen codebook_kind + bool force_random_rotation + + cdef cppclass index[IdxT](ann_index): + index(const handle_t& handle, + DistanceType metric, + codebook_gen codebook_kind, + uint32_t n_lists, + uint32_t dim, + uint32_t pq_bits, + uint32_t pq_dim, + uint32_t n_nonempty_lists) + + cdef cppclass search_params(ann_search_params): + uint32_t n_probes + cudaDataType_t lut_dtype + cudaDataType_t internal_distance_dtype + uint32_t preferred_thread_block_size + +cdef extern from "raft/neighbors/specializations/ivf_pq_specialization.hpp" \ + namespace "raft::neighbors::ivf_pq": + + cdef void search[T, IdxT](const handle_t& handle, + const search_params& params, + const index[IdxT]& index, + const T* queries, + uint32_t n_queries, + uint32_t k, + IdxT* neighbors, + float* distances, + device_memory_resource* mr) + + cdef index[int64_t] build(const handle_t& handle, + const index_params& params, + const float* dataset, + int64_t n_rows, + uint32_t dim) + + # cdef index[int64_t] build(const handle_t& handle, + # const index_params& params, + # const int8_t* dataset, + # int64_t n_rows, + # uint32_t dim) + + # cdef index[int64_t] build(const handle_t& handle, + # const index_params& params, + # const uint8_t* dataset, + # int64_t n_rows, + # uint32_t dim) + + cdef index[IdxT] extend[T, IdxT](const handle_t& handle, + const index[IdxT]& orig_index, + const T* new_vectors, + const IdxT* new_indices, + IdxT n_rows) + + +def is_c_cont(cai, dt): + return "strides" not in cai or \ + cai["strides"] is None or \ + cai["strides"][1] == dt.itemsize + + +def _get_codebook_kind(kind): + return { + 'per_subspace': 0, # codebook_gen.PER_SUBSPACE, + 'per_cluster': 1 #codebook_gen.PER_CLUSTER + }[kind] + + +def _get_metric(metric): + SUPPORTED_DISTANCES = { + "sqeuclidean": DistanceType.L2Expanded, + "euclidean": DistanceType.L2SqrtExpanded, + "inner_product": DistanceType.InnerProduct + } + if metric not in SUPPORTED_DISTANCES: + raise ValueError("metric %s is not supported" % metric) + return SUPPORTED_DISTANCES[metric] + + +class IvfPq: + """ + Nearest neighbors search using IVF-PQ method. + """ + + def __init__(self, handle=None): + self.handle = pylibraft.common.handle.Handle() if handle is None \ + else handle + self._index = None + + def __del__(self): + self._dealloc() + + def _dealloc(self): + # deallocate the index + cdef index[int64_t] *idx = self._index + del idx + + def build(self, + dataset, + n_lists = 1024, + metric="euclidean", + kmeans_n_iters=20, + kmeans_trainset_fraction=0.5, + pq_bits=8, + pq_dim=0, + codebook_kind="per_subspace", + force_random_rotation=0, + add_data_on_build=True): + + """ + Builds an IVF-PQ index that can be later used for nearest neighbor search. + + Parameters + ---------- + + dateset : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [float, int8, uint8] + n_list : int, default = 1024 + The number of clusters used in the coarse quantizer. + metric : string denoting the metric type, default="euclidean" + Valid values for metric: ["euclidean", "sqeuclidean", "inner_product"], + where sqeuclidean is the equclidean distance without the square root operation. + kmeans_trainset_fraction : int, default = 0.5 + If kmeans_trainset_fraction is less than 1, then the dataset is subsampled, + and only n_samples * kmeans_trainset_fraction rows are used for training. + pq_bits : int, default = 8 + The bit length of the vector element after quantization. + pq_dim : int, default = 0 + The dimensionality of a the vector after product quantization. When zero, an + optimal value is selected using a heuristic. Note pq_dim * pq_bits must be a multiple of 8. + Hint: a smaller 'pq_dim' results in a smaller index size and better search performance, but + lower recall. If 'pq_bits' is 8, 'pq_dim' can be set to any number, but multiple of 8 are + desirable for good performance. If 'pq_bits' is not 8, 'pq_dim' should be a multiple of 8. + For good performance, it is desirable that 'pq_dim' is a multiple of 32. Ideally, 'pq_dim' + should be also a divisor of the dataset dim. + codebook_kind : string, default = "per_subspace" + Valid values ["per_subspace", "per_cluster"] + force_random_rotation : bool, default = False + Apply a random rotation matrix on the input data and queries even if `dim % pq_dim == 0`. + Note: if `dim` is not multiple of `pq_dim`, a random rotation is always applied to the input + data and queries to transform the working space from `dim` to `rot_dim`, which may be slightly + larger than the original space and and is a multiple of `pq_dim` (`rot_dim % pq_dim == 0`). + However, this transform is not necessary when `dim` is multiple of `pq_dim` + (`dim == rot_dim`, hence no need in adding "extra" data columns / features). + By default, if `dim == rot_dim`, the rotation transform is initialized with the identity + matrix. When `force_random_rotation == True`, a random orthogonal transform matrix is generated + regardless of the values of `dim` and `pq_dim`. + add_data_on_build : bool, default = True + 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 + + import cupy as cp + + from pylibraft.neighbors import IvfPq + + 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), + 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) + """ + # TODO(tfeher): ensure that this works with managed memory as well + dataset_cai = dataset.__cuda_array_interface__ + dataset_dt = np.dtype(dataset_cai["typestr"]) + if not is_c_cont(dataset_cai, dataset_dt): + raise ValueError("Row major input is expected") + + cdef index_params params + params.n_lists = n_lists + params.metric = _get_metric(metric) + params.metric_arg = 0 + params.kmeans_n_iters = kmeans_n_iters + params.kmeans_trainset_fraction = kmeans_trainset_fraction + params.pq_bits = pq_bits + params.pq_dim = pq_dim + #params.codebook_kind = _get_codebook_kind(codebook_kind) + params.force_random_rotation = force_random_rotation + params.add_data_on_build = add_data_on_build + + + # cdef index[int64_t] *index_ptr + cdef int64_t n_rows = dataset_cai["shape"][0] # make it uint32_t + cdef uint32_t dim = dataset_cai["shape"][1] + cdef handle_t* handle_ = self.handle.getHandle() + cdef uintptr_t dataset_ptr = dataset_cai["data"][0] + + self._dealloc() + + cdef index[int64_t] *idx = new index[int64_t](deref(handle_), + _get_metric(metric), + codebook_gen.PER_SUBSPACE , + n_lists, + dim, + pq_bits, + pq_dim, + 0) + + if dataset_dt == np.float32: + idx[0] = build(deref(handle_), + params, + dataset_ptr, + n_rows, + dim) + self._index = idx + # elif dataset_dt == np.int8: + # idx[0] = build(deref(handle_), + # params, + # dataset_ptr, + # n_rows, + # dim) + # self._index = idx + # elif dataset_dt == np.uint8: + # idx[0] = build(deref(handle_), + # params, + # dataset_ptr, + # n_rows, + # dim) + else: + raise ValueError("dtype %s not supported" % dataset_dt) + + self.handle.sync() + + # def search(self, queries, k, neighbors=None, distances=None, n_probes=20, + # lut_dtype=CUDA_R_32F, internal_distance_type=CUDA_R_32F, + # preferred_thread_block_size=0): + # assert(numProbes > 0) + # assert(topK > 0) + # assert(numQueries > 0) + + # numQueries, dimQueries = queries.shape + # assert(dimQueries == self._index_params['dimDataset']) + # dtype = queries.dtype + + # if neighbors is None: + # neighbors = cupy.zeros( + # (numQueries, self._search_params['topK']), dtype='uint64') + # assert(isinstance(neighbors, cupy.ndarray)) + # assert(neighbors.shape[0] == numQueries) + # assert(neighbors.shape[1] == self._search_params['topK']) + # assert(neighbors.dtype is numpy.dtype('uint64')) + + # if distances is None: + # distances = cupy.zeros( + # (numQueries, self._search_params['topK']), dtype='float32') + # assert(isinstance(distances, cupy.ndarray)) + # assert(distances.shape[0] == numQueries) + # assert(distances.shape[1] == self._search_params['topK']) + # assert(distances.dtype is numpy.dtype('float32')) + + # search(handle_, + # params, + # index, + # queries_ptr, + # n_queries, + # k, + # neighbors_ptr, + # distances, + # nullptr) + + # return neighbors, distances + diff --git a/python/pylibraft/pylibraft/neighbors/neighbors.pyx b/python/pylibraft/pylibraft/neighbors/neighbors.pyx deleted file mode 100644 index 8e35087132..0000000000 --- a/python/pylibraft/pylibraft/neighbors/neighbors.pyx +++ /dev/null @@ -1,89 +0,0 @@ -# -# Copyright (c) 2022, NVIDIA CORPORATION. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# cython: profile=False -# distutils: language = c++ -# cython: embedsignature = True -# cython: language_level = 3 - -import numpy as np - -from libc.stdint cimport uintptr_t -from libc.stdint cimport uint32_t -from cython.operator cimport dereference as deref -from cuda.ccudart cimport cudaDataType_t - -from libcpp cimport bool -from pylibraft.distance.distance_type cimport DistanceType -from pylibraft.common.handle cimport handle_t -from rmm._lib.memory_resource cimport device_memory_resource - -def is_c_cont(cai, dt): - return "strides" not in cai or \ - cai["strides"] is None or \ - cai["strides"][1] == dt.itemsize - -cdef extern from "raft/neighbors/ann_types.hpp" \ - namespace "raft::neighbors::ann": - - cdef cppclass ann_index "raft::neighbors::index" - - cdef cppclass ann_index_params "raft::spatial::knn::index_params": - DistanceType metric - float metric_arg - bool add_data_on_build - - cdef cppclass ann_search_params "raft::spatial::knn::search_params": - pass - -cdef extern from "raft/neighbors/ivf_pq_types.hpp" \ - namespace "raft::neighbors::ivf_pq": - - cdef enum codebook_gen: - PER_SUBSPACE = 0, - PER_CLUSTER = 1, - - - cdef cppclass index_params(ann_index_params): - uint32_t n_lists - uint32_t kmeans_n_iters - double kmeans_trainset_fraction - uint32_t pq_bits - uint32_t pq_dim - codebook_gen codebook_kind - bool force_random_rotation - - cdef cppclass index[IdxT](ann_index) - - cdef cppclass search_params(ann_search_params): - uint32_t n_probes - cudaDataType_t lut_dtype - cudaDataType_t internal_distance_dtype - uint32_t preferred_thread_block_size - -cdef extern from "raft/neighbors/specializations/ivf_pq_specialization.hpp" \ - namespace "raft::neighbors ::ivf_pq": - - cdef void search[T, IdxT](const handle_t& handle, - const search_params& params, - const index[IdxT]& index, - const T* queries, - uint32_t n_queries, - uint32_t k, - IdxT* neighbors, - float* distances, - device_memory_resource* mr) - - # add build specialization and wrapper, use them \ No newline at end of file