diff --git a/.gitignore b/.gitignore index 68996dbdf..4b6f46320 100644 --- a/.gitignore +++ b/.gitignore @@ -26,6 +26,8 @@ dask-worker-space/ *.bin bench/ann/data temporary_*.json +rust/target/ +rust/Cargo.lock ## scikit-build _skbuild diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 576b53907..b1de648b0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -114,7 +114,6 @@ repos: setup[.]cfg$ exclude: | (?x) - cpp/include/cuvs/neighbors/detail/faiss_select/| docs/source/sphinxext/github_link\.py| cpp/cmake/modules/FindAVX\.cmake| diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index b8f15aba7..0082461c5 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -37,7 +37,7 @@ function sed_runner() { sed -i.bak ''"$1"'' $2 && rm -f ${2}.bak } -sed_runner "s/set(RAPIDS_VERSION .*)/set(RAPIDS_VERSION \"${NEXT_SHORT_TAG}\")/g" cpp/examples/cmake/thirdparty/fetch_rapids.cmake +sed_runner "s/set(RAPIDS_VERSION .*)/set(RAPIDS_VERSION \"${NEXT_SHORT_TAG}\")/g" examples/cmake/thirdparty/fetch_rapids.cmake # Centralized version file update echo "${NEXT_FULL_TAG}" > VERSION @@ -47,6 +47,7 @@ DEPENDENCIES=( cuvs cuvs-cu11 cuvs-cu12 + pylibraft pylibraft-cu11 pylibraft-cu12 rmm @@ -68,19 +69,13 @@ for FILE in python/*/pyproject.toml; do sed_runner "/\"ucx-py==/ s/==.*\"/==${NEXT_UCX_PY_SHORT_TAG_PEP440}.*\"/g" ${FILE} done -sed_runner "/^ucx_py_version:$/ {n;s/.*/ - \"${NEXT_UCX_PY_VERSION}\"/}" conda/recipes/raft-dask/conda_build_config.yaml - for FILE in .github/workflows/*.yaml; do sed_runner "/shared-workflows/ s/@.*/@branch-${NEXT_SHORT_TAG}/g" "${FILE}" done -sed_runner "/^set(CUVS_VERSION/ s|\".*\"|\"${NEXT_SHORT_TAG}\"|g" docs/source/build.md -sed_runner "s|branch-[0-9][0-9].[0-9][0-9]|branch-${NEXT_SHORT_TAG}|g" docs/source/build.md sed_runner "/rapidsai\/raft/ s|branch-[0-9][0-9].[0-9][0-9]|branch-${NEXT_SHORT_TAG}|g" docs/source/developer_guide.md -sed_runner "s|:[0-9][0-9].[0-9][0-9]|:${NEXT_SHORT_TAG}|g" docs/source/raft_ann_benchmarks.md - -sed_runner "s|branch-[0-9][0-9].[0-9][0-9]|branch-${NEXT_SHORT_TAG}|g" README.md +sed_runner "s|=[0-9][0-9].[0-9][0-9]|=${NEXT_SHORT_TAG}|g" README.md # rust can't handle leading 0's in the major/minor/patch version - remove NEXT_FULL_RUST_TAG=$(printf "%d.%d.%d" $((10#$NEXT_MAJOR)) $((10#$NEXT_MINOR)) $((10#$NEXT_PATCH))) @@ -90,6 +85,7 @@ sed_runner "s/version = \".*\"/version = \"${NEXT_FULL_RUST_TAG}\"/g" rust/Cargo find .devcontainer/ -type f -name devcontainer.json -print0 | while IFS= read -r -d '' filename; do sed_runner "s@rapidsai/devcontainers:[0-9.]*@rapidsai/devcontainers:${NEXT_SHORT_TAG}@g" "${filename}" sed_runner "s@rapidsai/devcontainers/features/ucx:[0-9.]*@rapidsai/devcontainers/features/ucx:${NEXT_SHORT_TAG_PEP440}@" "${filename}" + sed_runner "s@rapidsai/devcontainers/features/cuda:[0-9.]*@rapidsai/devcontainers/features/cuda:${NEXT_SHORT_TAG_PEP440}@" "${filename}" sed_runner "s@rapidsai/devcontainers/features/rapids-build-utils:[0-9.]*@rapidsai/devcontainers/features/rapids-build-utils:${NEXT_SHORT_TAG_PEP440}@" "${filename}" sed_runner "s@rapids-\${localWorkspaceFolderBasename}-${CURRENT_SHORT_TAG}@rapids-\${localWorkspaceFolderBasename}-${NEXT_SHORT_TAG}@g" "${filename}" done diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5715ee8a2..b972fc122 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -447,14 +447,29 @@ target_include_directories( "$" ) +rapids_find_package( + OpenMP REQUIRED + BUILD_EXPORT_SET cuvs-exports + INSTALL_EXPORT_SET cuvs-exports +) + if(NOT BUILD_CPU_ONLY) + + set(CUVS_CUSOLVER_DEPENDENCY CUDA::cusolver${_ctk_static_suffix}) + set(CUVS_CUBLAS_DEPENDENCY CUDA::cublas${_ctk_static_suffix}) + set(CUVS_CURAND_DEPENDENCY CUDA::curand${_ctk_static_suffix}) + set(CUVS_CUSPARSE_DEPENDENCY CUDA::cusparse${_ctk_static_suffix}) + + set(CUVS_CTK_MATH_DEPENDENCIES ${CUVS_CUBLAS_DEPENDENCY} ${CUVS_CUSOLVER_DEPENDENCY} + ${CUVS_CUSPARSE_DEPENDENCY} ${CUVS_CURAND_DEPENDENCY} + ) + # Keep cuVS as lightweight as possible. Only CUDA libs and rmm should be used in global target. target_link_libraries( cuvs - PUBLIC rmm::rmm $<$>:raft::raft> - $<$>:raft::compiled> - PRIVATE $<$:raft::raft> - $<$:raft::compiled_static> nvidia::cutlass::cutlass + PUBLIC rmm::rmm raft::raft + PRIVATE nvidia::cutlass::cutlass ${CUVS_CTK_MATH_DEPENDENCIES} + $ ) endif() @@ -559,7 +574,7 @@ if(BUILD_C_LIBRARY) target_link_libraries( cuvs_c PUBLIC cuvs::cuvs - PRIVATE $<$:raft::raft> + PRIVATE raft::raft ${CUVS_CTK_MATH_DEPENDENCIES} ) # ensure CUDA symbols aren't relocated to the middle of the debug build binaries diff --git a/cpp/cmake/thirdparty/get_raft.cmake b/cpp/cmake/thirdparty/get_raft.cmake index e3c2de320..44825f91c 100644 --- a/cpp/cmake/thirdparty/get_raft.cmake +++ b/cpp/cmake/thirdparty/get_raft.cmake @@ -17,7 +17,7 @@ set(RAFT_FORK "rapidsai") set(RAFT_PINNED_TAG "branch-${RAPIDS_VERSION_MAJOR_MINOR}") function(find_and_configure_raft) - set(oneValueArgs VERSION FORK PINNED_TAG COMPILE_LIBRARY USE_RAFT_STATIC ENABLE_NVTX ENABLE_MNMG_DEPENDENCIES) + set(oneValueArgs VERSION FORK PINNED_TAG USE_RAFT_STATIC ENABLE_NVTX ENABLE_MNMG_DEPENDENCIES) cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN} ) @@ -31,16 +31,6 @@ function(find_and_configure_raft) set(RAFT_COMPONENTS "") - if(PKG_COMPILE_LIBRARY) - if(NOT PKG_USE_RAFT_STATIC) - string(APPEND RAFT_COMPONENTS " compiled") - set(RAFT_COMPILED_LIB raft::compiled PARENT_SCOPE) - else() - string(APPEND RAFT_COMPONENTS " compiled_static") - set(RAFT_COMPILED_LIB raft::compiled_static PARENT_SCOPE) - endif() - endif() - if(PKG_ENABLE_MNMG_DEPENDENCIES) string(APPEND RAFT_COMPONENTS " distributed") endif() @@ -62,7 +52,7 @@ function(find_and_configure_raft) "BUILD_PRIMS_BENCH OFF" "BUILD_ANN_BENCH OFF" "RAFT_NVTX ${PKG_ENABLE_NVTX}" - "RAFT_COMPILE_LIBRARY ${PKG_COMPILE_LIBRARY}" + "RAFT_COMPILE_LIBRARY OFF" ) endfunction() @@ -72,7 +62,6 @@ endfunction() find_and_configure_raft(VERSION ${RAFT_VERSION}.00 FORK ${RAFT_FORK} PINNED_TAG ${RAFT_PINNED_TAG} - COMPILE_LIBRARY ON ENABLE_MNMG_DEPENDENCIES OFF ENABLE_NVTX OFF USE_RAFT_STATIC ${CUVS_USE_RAFT_STATIC} diff --git a/cpp/include/cuvs/cluster/agglomerative.hpp b/cpp/include/cuvs/cluster/agglomerative.hpp index 5da03c4d1..658f5628c 100644 --- a/cpp/include/cuvs/cluster/agglomerative.hpp +++ b/cpp/include/cuvs/cluster/agglomerative.hpp @@ -16,7 +16,7 @@ #pragma once -#include +#include #include #include #include diff --git a/cpp/include/cuvs/cluster/kmeans.hpp b/cpp/include/cuvs/cluster/kmeans.hpp index 471a4d12d..7a372b1bc 100644 --- a/cpp/include/cuvs/cluster/kmeans.hpp +++ b/cpp/include/cuvs/cluster/kmeans.hpp @@ -14,7 +14,7 @@ * limitations under the License. */ #pragma once -#include +#include #include #include #include diff --git a/cpp/include/cuvs/distance/distance_types.h b/cpp/include/cuvs/distance/distance.h similarity index 100% rename from cpp/include/cuvs/distance/distance_types.h rename to cpp/include/cuvs/distance/distance.h diff --git a/cpp/include/cuvs/distance/pairwise_distance.hpp b/cpp/include/cuvs/distance/distance.hpp similarity index 81% rename from cpp/include/cuvs/distance/pairwise_distance.hpp rename to cpp/include/cuvs/distance/distance.hpp index 5f774e3bc..5786b0a32 100644 --- a/cpp/include/cuvs/distance/pairwise_distance.hpp +++ b/cpp/include/cuvs/distance/distance.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2021-2023, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -16,13 +16,54 @@ #pragma once +#include "distance.h" + #include -#include #include #include namespace cuvs::distance { +using DistanceType = cuvsDistanceType; + +/** + * Whether minimal distance corresponds to similar elements (using the given metric). + */ +inline bool is_min_close(DistanceType metric) +{ + bool select_min; + switch (metric) { + case DistanceType::InnerProduct: + // Similarity metrics have the opposite meaning, i.e. nearest neighbors are those with larger + // similarity (See the same logic at cpp/include/raft/sparse/spatial/detail/knn.cuh:362 + // {perform_k_selection}) + select_min = false; + break; + default: select_min = true; + } + return select_min; +} + +namespace kernels { +enum KernelType { LINEAR, POLYNOMIAL, RBF, TANH }; + +/** + * Parameters for kernel matrices. + * The following kernels are implemented: + * - LINEAR \f[ K(x_1,x_2) = , \f] where \f$< , >\f$ is the dot product + * - POLYNOMIAL \f[ K(x_1, x_2) = (\gamma + \mathrm{coef0})^\mathrm{degree} \f] + * - RBF \f[ K(x_1, x_2) = \exp(- \gamma |x_1-x_2|^2) \f] + * - TANH \f[ K(x_1, x_2) = \tanh(\gamma + \mathrm{coef0}) \f] + */ +struct KernelParams { + // Kernel function parameters + KernelType kernel; //!< Type of the kernel function + int degree; //!< Degree of polynomial kernel (ignored by others) + double gamma; //!< multiplier in the + double coef0; //!< additive constant in poly and tanh kernels +}; +} // end namespace kernels + /** * @defgroup pairwise_distance Pairwise Distances API * @{ @@ -37,7 +78,7 @@ namespace cuvs::distance { * @code{.cpp} * #include * #include - * #include + * #include * * raft::resources handle; * int n_samples = 5000; @@ -81,7 +122,7 @@ void pairwise_distance( * @code{.cpp} * #include * #include - * #include + * #include * * raft::resources handle; * int n_samples = 5000; @@ -125,7 +166,7 @@ void pairwise_distance( * @code{.cpp} * #include * #include - * #include + * #include * * raft::resources handle; * int n_samples = 5000; @@ -168,7 +209,7 @@ void pairwise_distance( * @code{.cpp} * #include * #include - * #include + * #include * * raft::resources handle; * int n_samples = 5000; @@ -205,4 +246,4 @@ void pairwise_distance( /** @} */ // end group pairwise_distance_runtime -} // namespace cuvs::distance +}; // namespace cuvs::distance diff --git a/cpp/include/cuvs/distance/distance_types.hpp b/cpp/include/cuvs/distance/distance_types.hpp deleted file mode 100644 index 8d6eb7137..000000000 --- a/cpp/include/cuvs/distance/distance_types.hpp +++ /dev/null @@ -1,63 +0,0 @@ -/* - * Copyright (c) 2021-2023, 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. - */ - -#pragma once - -#include "distance_types.h" - -namespace cuvs::distance { - -using DistanceType = cuvsDistanceType; - -/** - * Whether minimal distance corresponds to similar elements (using the given metric). - */ -inline bool is_min_close(DistanceType metric) -{ - bool select_min; - switch (metric) { - case DistanceType::InnerProduct: - // Similarity metrics have the opposite meaning, i.e. nearest neighbors are those with larger - // similarity (See the same logic at cpp/include/raft/sparse/spatial/detail/knn.cuh:362 - // {perform_k_selection}) - select_min = false; - break; - default: select_min = true; - } - return select_min; -} - -namespace kernels { -enum KernelType { LINEAR, POLYNOMIAL, RBF, TANH }; - -/** - * Parameters for kernel matrices. - * The following kernels are implemented: - * - LINEAR \f[ K(x_1,x_2) = , \f] where \f$< , >\f$ is the dot product - * - POLYNOMIAL \f[ K(x_1, x_2) = (\gamma + \mathrm{coef0})^\mathrm{degree} \f] - * - RBF \f[ K(x_1, x_2) = \exp(- \gamma |x_1-x_2|^2) \f] - * - TANH \f[ K(x_1, x_2) = \tanh(\gamma + \mathrm{coef0}) \f] - */ -struct KernelParams { - // Kernel function parameters - KernelType kernel; //!< Type of the kernel function - int degree; //!< Degree of polynomial kernel (ignored by others) - double gamma; //!< multiplier in the - double coef0; //!< additive constant in poly and tanh kernels -}; -} // end namespace kernels - -}; // namespace cuvs::distance diff --git a/cpp/include/cuvs/neighbors/ann_types.hpp b/cpp/include/cuvs/neighbors/ann_types.hpp deleted file mode 100644 index e9b7bd903..000000000 --- a/cpp/include/cuvs/neighbors/ann_types.hpp +++ /dev/null @@ -1,90 +0,0 @@ -/* - * Copyright (c) 2022-2023, 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. - */ - -#pragma once - -#include -#include - -namespace cuvs::neighbors::ann { - -/** Parameters for VPQ compression. */ -struct vpq_params { - /** - * The bit length of the vector element after compression by PQ. - * - * Possible values: [4, 5, 6, 7, 8]. - * - * Hint: the smaller the 'pq_bits', the smaller the index size and the better the search - * performance, but the lower the recall. - */ - uint32_t pq_bits = 8; - /** - * The dimensionality of the vector after compression by PQ. - * When zero, an optimal value is selected using a heuristic. - * - * TODO: at the moment `dim` must be a multiple `pq_dim`. - */ - uint32_t pq_dim = 0; - /** - * Vector Quantization (VQ) codebook size - number of "coarse cluster centers". - * When zero, an optimal value is selected using a heuristic. - */ - uint32_t vq_n_centers = 0; - /** The number of iterations searching for kmeans centers (both VQ & PQ phases). */ - uint32_t kmeans_n_iters = 25; - /** - * The fraction of data to use during iterative kmeans building (VQ phase). - * When zero, an optimal value is selected using a heuristic. - */ - double vq_kmeans_trainset_fraction = 0; - /** - * The fraction of data to use during iterative kmeans building (PQ phase). - * When zero, an optimal value is selected using a heuristic. - */ - double pq_kmeans_trainset_fraction = 0; -}; - -/** - * @defgroup ann_types Approximate Nearest Neighbors Types - * @{ - */ - -/** The base for approximate KNN index structures. */ -struct index {}; - -/** The base for KNN index parameters. */ -struct index_params { - /** Distance type. */ - cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded; - /** The argument used by some distance metrics. */ - float metric_arg = 2.0f; - /** - * Whether to add the dataset content to the index, i.e.: - * - * - `true` means the index is filled with the dataset vectors and ready to search after calling - * `build`. - * - `false` means `build` only trains the underlying model (e.g. quantizer or clustering), but - * the index is left empty; you'd need to call `extend` on the index afterwards to populate it. - */ - bool add_data_on_build = true; -}; - -struct search_params {}; - -/** @} */ // end group ann_types - -}; // namespace cuvs::neighbors::ann diff --git a/cpp/include/cuvs/neighbors/brute_force.h b/cpp/include/cuvs/neighbors/brute_force.h index 145bb5555..e285eae37 100644 --- a/cpp/include/cuvs/neighbors/brute_force.h +++ b/cpp/include/cuvs/neighbors/brute_force.h @@ -17,7 +17,7 @@ #pragma once #include -#include +#include #include #include diff --git a/cpp/include/cuvs/neighbors/brute_force.hpp b/cpp/include/cuvs/neighbors/brute_force.hpp index 755a94122..755c8cfdb 100644 --- a/cpp/include/cuvs/neighbors/brute_force.hpp +++ b/cpp/include/cuvs/neighbors/brute_force.hpp @@ -16,8 +16,8 @@ #pragma once -#include "ann_types.hpp" -#include +#include "common.hpp" +#include #include #include #include @@ -37,7 +37,7 @@ namespace cuvs::neighbors::brute_force { * @tparam T data element type */ template -struct index : cuvs::neighbors::ann::index { +struct index : cuvs::neighbors::index { public: index(const index&) = delete; index(index&&) = default; diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index e421c9832..fd7388211 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -16,9 +16,9 @@ #pragma once -#include "ann_types.hpp" -#include -#include +#include "common.hpp" +#include +#include #include #include #include @@ -49,7 +49,7 @@ enum class graph_build_algo { NN_DESCENT }; -struct index_params : ann::index_params { +struct index_params : cuvs::neighbors::index_params { /** Degree of input graph for pruning. */ size_t intermediate_graph_degree = 128; /** Degree of output graph. */ @@ -63,7 +63,7 @@ struct index_params : ann::index_params { * * NOTE: this is experimental new API, consider it unsafe. */ - std::optional compression = std::nullopt; + std::optional compression = std::nullopt; }; /** @@ -86,7 +86,7 @@ enum class search_algo { enum class hash_mode { HASH, SMALL, AUTO }; -struct search_params : ann::search_params { +struct search_params : cuvs::neighbors::search_params { /** Maximum number of queries to search at the same time (batch size). Auto select when 0.*/ size_t max_queries = 0; @@ -151,7 +151,7 @@ static_assert(std::is_aggregate_v); * */ template -struct index : ann::index { +struct index : cuvs::neighbors::index { static_assert(!raft::is_narrowing_v, "IdxT must be able to represent all values of uint32_t"); @@ -209,7 +209,7 @@ struct index : ann::index { /** Construct an empty index. */ index(raft::resources const& res, cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded) - : ann::index(), + : cuvs::neighbors::index(), metric_(metric), graph_(raft::make_device_matrix(res, 0, 0)), dataset_(new cuvs::neighbors::empty_dataset(0)) @@ -276,7 +276,7 @@ struct index : ann::index { raft::mdspan, raft::row_major, data_accessor> dataset, raft::mdspan, raft::row_major, graph_accessor> knn_graph) - : ann::index(), + : cuvs::neighbors::index(), metric_(metric), graph_(raft::make_device_matrix(res, 0, 0)), dataset_(make_aligned_dataset(res, dataset, 16)) diff --git a/cpp/include/cuvs/neighbors/dataset.hpp b/cpp/include/cuvs/neighbors/common.hpp similarity index 53% rename from cpp/include/cuvs/neighbors/dataset.hpp rename to cpp/include/cuvs/neighbors/common.hpp index 79c3f644f..45fa1a107 100644 --- a/cpp/include/cuvs/neighbors/dataset.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -13,9 +13,11 @@ * See the License for the specific language governing permissions and * limitations under the License. */ + #pragma once -#include +#include +#include #include #include #include @@ -23,6 +25,9 @@ #include // get_device_for_address #include // rounding up +#include +#include + #include #include #include @@ -33,6 +38,72 @@ namespace cuvs::neighbors { +/** Parameters for VPQ compression. */ +struct vpq_params { + /** + * The bit length of the vector element after compression by PQ. + * + * Possible values: [4, 5, 6, 7, 8]. + * + * Hint: the smaller the 'pq_bits', the smaller the index size and the better the search + * performance, but the lower the recall. + */ + uint32_t pq_bits = 8; + /** + * The dimensionality of the vector after compression by PQ. + * When zero, an optimal value is selected using a heuristic. + * + * TODO: at the moment `dim` must be a multiple `pq_dim`. + */ + uint32_t pq_dim = 0; + /** + * Vector Quantization (VQ) codebook size - number of "coarse cluster centers". + * When zero, an optimal value is selected using a heuristic. + */ + uint32_t vq_n_centers = 0; + /** The number of iterations searching for kmeans centers (both VQ & PQ phases). */ + uint32_t kmeans_n_iters = 25; + /** + * The fraction of data to use during iterative kmeans building (VQ phase). + * When zero, an optimal value is selected using a heuristic. + */ + double vq_kmeans_trainset_fraction = 0; + /** + * The fraction of data to use during iterative kmeans building (PQ phase). + * When zero, an optimal value is selected using a heuristic. + */ + double pq_kmeans_trainset_fraction = 0; +}; + +/** + * @defgroup neighbors_index Approximate Nearest Neighbors Types + * @{ + */ + +/** The base for approximate KNN index structures. */ +struct index {}; + +/** The base for KNN index parameters. */ +struct index_params { + /** Distance type. */ + cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded; + /** The argument used by some distance metrics. */ + float metric_arg = 2.0f; + /** + * Whether to add the dataset content to the index, i.e.: + * + * - `true` means the index is filled with the dataset vectors and ready to search after calling + * `build`. + * - `false` means `build` only trains the underlying model (e.g. quantizer or clustering), but + * the index is left empty; you'd need to call `extend` on the index afterwards to populate it. + */ + bool add_data_on_build = true; +}; + +struct search_params {}; + +/** @} */ // end group neighbors_index + /** Two-dimensional dataset; maybe owning, maybe compressed, maybe strided. */ template struct dataset { @@ -207,10 +278,6 @@ auto make_aligned_dataset(const raft::resources& res, const SrcT& src, uint32_t raft::round_up_safe(src.extent(1) * kSize, std::lcm(align_bytes, kSize)) / kSize; return make_strided_dataset(res, src, required_stride); } - -/** Parameters for VPQ compression. */ -using vpq_params = cuvs::neighbors::ann::vpq_params; - /** * @brief VPQ compressed dataset. * @@ -295,4 +362,240 @@ struct vpq_dataset : public dataset { } }; -} // namespace cuvs::neighbors +namespace filtering { + +/* A filter that filters nothing. This is the default behavior. */ +struct none_ivf_sample_filter { + inline _RAFT_HOST_DEVICE bool operator()( + // query index + const uint32_t query_ix, + // the current inverted list index + const uint32_t cluster_ix, + // the index of the current sample inside the current inverted list + const uint32_t sample_ix) const; +}; + +/* A filter that filters nothing. This is the default behavior. */ +struct none_cagra_sample_filter { + inline _RAFT_HOST_DEVICE bool operator()( + // query index + const uint32_t query_ix, + // the index of the current sample + const uint32_t sample_ix) const; +}; + +/** + * @brief Filter used to convert the cluster index and sample index + * of an IVF search into a sample index. This can be used as an + * intermediate filter. + * + * @tparam index_t Indexing type + * @tparam filter_t + */ +template +struct ivf_to_sample_filter { + const index_t* const* inds_ptrs_; + const filter_t next_filter_; + + ivf_to_sample_filter(const index_t* const* inds_ptrs, const filter_t next_filter); + + /** If the original filter takes three arguments, then don't modify the arguments. + * If the original filter takes two arguments, then we are using `inds_ptr_` to obtain the sample + * index. + */ + inline _RAFT_HOST_DEVICE bool operator()( + // query index + const uint32_t query_ix, + // the current inverted list index + const uint32_t cluster_ix, + // the index of the current sample inside the current inverted list + const uint32_t sample_ix) const; +}; + +/** + * @brief Filter an index with a bitset + * + * @tparam index_t Indexing type + */ +template +struct bitset_filter { + // View of the bitset to use as a filter + const cuvs::core::bitset_view bitset_view_; + + bitset_filter(const cuvs::core::bitset_view bitset_for_filtering); + inline _RAFT_HOST_DEVICE bool operator()( + // query index + const uint32_t query_ix, + // the index of the current sample + const uint32_t sample_ix) const; +}; + +/** + * If the filtering depends on the index of a sample, then the following + * filter template can be used: + * + * template + * struct index_ivf_sample_filter { + * using index_type = IdxT; + * + * const index_type* const* inds_ptr = nullptr; + * + * index_ivf_sample_filter() {} + * index_ivf_sample_filter(const index_type* const* _inds_ptr) + * : inds_ptr{_inds_ptr} {} + * index_ivf_sample_filter(const index_ivf_sample_filter&) = default; + * index_ivf_sample_filter(index_ivf_sample_filter&&) = default; + * index_ivf_sample_filter& operator=(const index_ivf_sample_filter&) = default; + * index_ivf_sample_filter& operator=(index_ivf_sample_filter&&) = default; + * + * inline _RAFT_HOST_DEVICE bool operator()( + * const uint32_t query_ix, + * const uint32_t cluster_ix, + * const uint32_t sample_ix) const { + * index_type database_idx = inds_ptr[cluster_ix][sample_ix]; + * + * // return true or false, depending on the database_idx + * return true; + * } + * }; + * + * Initialize it as: + * using filter_type = index_ivf_sample_filter; + * filter_type filter(cuvs_ivfpq_index.inds_ptrs().data_handle()); + * + * Use it as: + * cuvs::neighbors::ivf_pq::search_with_filtering( + * ...regular parameters here..., + * filter + * ); + * + * Another example would be the following filter that greenlights samples according + * to a contiguous bit mask vector. + * + * template + * struct bitmask_ivf_sample_filter { + * using index_type = IdxT; + * + * const index_type* const* inds_ptr = nullptr; + * const uint64_t* const bit_mask_ptr = nullptr; + * const int64_t bit_mask_stride_64 = 0; + * + * bitmask_ivf_sample_filter() {} + * bitmask_ivf_sample_filter( + * const index_type* const* _inds_ptr, + * const uint64_t* const _bit_mask_ptr, + * const int64_t _bit_mask_stride_64) + * : inds_ptr{_inds_ptr}, + * bit_mask_ptr{_bit_mask_ptr}, + * bit_mask_stride_64{_bit_mask_stride_64} {} + * bitmask_ivf_sample_filter(const bitmask_ivf_sample_filter&) = default; + * bitmask_ivf_sample_filter(bitmask_ivf_sample_filter&&) = default; + * bitmask_ivf_sample_filter& operator=(const bitmask_ivf_sample_filter&) = default; + * bitmask_ivf_sample_filter& operator=(bitmask_ivf_sample_filter&&) = default; + * + * inline _RAFT_HOST_DEVICE bool operator()( + * const uint32_t query_ix, + * const uint32_t cluster_ix, + * const uint32_t sample_ix) const { + * const index_type database_idx = inds_ptr[cluster_ix][sample_ix]; + * const uint64_t bit_mask_element = + * bit_mask_ptr[query_ix * bit_mask_stride_64 + database_idx / 64]; + * const uint64_t masked_bool = + * bit_mask_element & (1ULL << (uint64_t)(database_idx % 64)); + * const bool is_bit_set = (masked_bool != 0); + * + * return is_bit_set; + * } + * }; + */ +} // namespace filtering + +namespace ivf { + +/** + * Default value filled in the `indices` array. + * One may encounter it trying to access a record within a list that is outside of the + * `size` bound or whenever the list is allocated but not filled-in yet. + */ +template +constexpr static IdxT kInvalidRecord = + (std::is_signed_v ? IdxT{0} : std::numeric_limits::max()) - 1; + +/** The data for a single IVF list. */ +template