From 04fa426230788f8205217894b458526d0644dc4d Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Thu, 30 Nov 2023 10:23:05 -0800 Subject: [PATCH 1/3] Add usage example for brute_force::build (#2029) Add a usage example for using the brute_force index api for building and searching. Also fix some minor compile time errors in the vector search tutorial Authors: - Ben Frederickson (https://github.com/benfred) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/raft/pull/2029 --- .../raft/neighbors/brute_force-inl.cuh | 40 +++++++++++++++++++ cpp/include/raft/neighbors/brute_force.cuh | 6 +-- docs/source/vector_search_tutorial.md | 28 ++++++------- 3 files changed, 57 insertions(+), 17 deletions(-) diff --git a/cpp/include/raft/neighbors/brute_force-inl.cuh b/cpp/include/raft/neighbors/brute_force-inl.cuh index 6b86f2463f..906371bd01 100644 --- a/cpp/include/raft/neighbors/brute_force-inl.cuh +++ b/cpp/include/raft/neighbors/brute_force-inl.cuh @@ -284,6 +284,44 @@ void fused_l2_knn(raft::resources const& handle, /** * @brief Build the index from the dataset for efficient search. * + * This function builds a brute force index for the given dataset. This lets you re-use + * precalculated norms for the dataset, leading to a speedup over calling + * raft::neighbors::brute_force::knn repeatedly. + * + * Example usage: + * @code{.cpp} + * #include + * #include + * #include + * + * // create a random dataset + * int n_rows = 10000; + * int n_cols = 10000; + * + * raft::device_resources res; + * auto dataset = raft::make_device_matrix(res, n_rows, n_cols); + * auto labels = raft::make_device_vector(res, n_rows); + * + * raft::random::make_blobs(res, dataset.view(), labels.view()); + * + * // create a brute_force knn index from the dataset + * auto index = raft::neighbors::brute_force::build(res, + * raft::make_const_mdspan(dataset.view())); + * + * // Use the constructed index to search for the nearest 128 neighbors + * int k = 128; + * auto search = raft::make_const_mdspan(dataset.view()); + * + * auto indices= raft::make_device_matrix(res, search.extent(0), k); + * auto distances = raft::make_device_matrix(res, search.extent(0), k); + * + * raft::neighbors::brute_force::search(res, + * index, + * search, + * indices.view(), + * distances.view()); + * @endcode + * * @tparam T data element type * * @param[in] res @@ -330,6 +368,8 @@ index build(raft::resources const& res, /** * @brief Brute Force search using the constructed index. * + * See raft::neighbors::brute_force::build for a usage example + * * @tparam T data element type * @tparam IdxT type of the indices * diff --git a/cpp/include/raft/neighbors/brute_force.cuh b/cpp/include/raft/neighbors/brute_force.cuh index 4ba9159556..331ea55540 100644 --- a/cpp/include/raft/neighbors/brute_force.cuh +++ b/cpp/include/raft/neighbors/brute_force.cuh @@ -44,10 +44,10 @@ namespace raft::neighbors::brute_force { * int n_cols = 10000; * raft::device_resources res; - * auto dataset = raft::make_device_matrix(res, n_rows, n_cols); - * auto labels = raft::make_device_vector(res, n_rows); + * auto dataset = raft::make_device_matrix(res, n_rows, n_cols); + * auto labels = raft::make_device_vector(res, n_rows); - * raft::make_blobs(res, dataset.view(), labels.view()); + * raft::random::make_blobs(res, dataset.view(), labels.view()); * * // create a brute_force knn index from the dataset * auto index = raft::neighbors::brute_force::build(res, diff --git a/docs/source/vector_search_tutorial.md b/docs/source/vector_search_tutorial.md index 126ac534c6..8ff25143e3 100644 --- a/docs/source/vector_search_tutorial.md +++ b/docs/source/vector_search_tutorial.md @@ -89,10 +89,10 @@ raft::device_resources res; int n_rows = 10000; int n_cols = 10000; -auto dataset = raft::make_device_matrix(res, n_rows, n_cols); -auto labels = raft::make_device_vector(res, n_rows); +auto dataset = raft::make_device_matrix(res, n_rows, n_cols); +auto labels = raft::make_device_vector(res, n_rows); -raft::make_blobs(res, dataset.view(), labels.view()); +raft::random::make_blobs(res, dataset.view(), labels.view()); ``` That's it. We've now generated a random 10kx10k matrix with points that cleanly separate into Gaussian clusters, along with a vector of cluster labels for each of the data points. Notice the `cuh` extension in the header file include for `make_blobs`. This signifies to us that this file contains CUDA device functions like kernel code so the CUDA compiler, `nvcc` is needed in order to compile any code that uses it. Generally, any source files that include headers with a `cuh` extension use the `.cu` extension instead of `.cpp`. The rule here is that `cpp` source files contain code which can be compiled with a C++ compiler like `g++` while `cu` files require the CUDA compiler. @@ -125,14 +125,14 @@ auto search = raft::make_const_mdspan(dataset.view()); // Indices and Distances are of dimensions (n, k) // where n is number of rows in the search matrix -auto reference_indices = raft::make_device_matrix(search.extent(0), k); // stores index of neighbors -auto reference_distances = raft::make_device_matrix(search.extent(0), k); // stores distance to neighbors +auto reference_indices = raft::make_device_matrix(res, search.extent(0), k); // stores index of neighbors +auto reference_distances = raft::make_device_matrix(res, search.extent(0), k); // stores distance to neighbors raft::neighbors::brute_force::search(res, bfknn_index, search, - raft::make_const_mdspan(indices.view()), - raft::make_const_mdspan(distances.view())); + reference_indices.view(), + reference_distances.view()); ``` We have established several things here by building a flat index. Now we know the exact 64 neighbors of all points in the matrix, and this algorithm can be generally useful in several ways: @@ -152,9 +152,9 @@ Next we'll train an ANN index. We'll use our graph-based CAGRA algorithm for thi raft::device_resources res; // use default index parameters -cagra::index_params index_params; +raft::neighbors::cagra::index_params index_params; -auto index = cagra::build(res, index_params, dataset); +auto index = raft::neighbors::cagra::build(res, index_params, raft::make_const_mdspan(dataset.view())); ``` ### Query the CAGRA index @@ -167,10 +167,10 @@ auto indices = raft::make_device_matrix(res, n_rows, k); auto distances = raft::make_device_matrix(res, n_rows, k); // use default search parameters -cagra::search_params search_params; +raft::neighbors::cagra::search_params search_params; // search K nearest neighbors -cagra::search( +raft::neighbors::cagra::search( res, search_params, index, search, indices.view(), distances.view()); ``` @@ -197,8 +197,8 @@ raft::stats::neighborhood_recall(res, raft::make_const_mdspan(indices.view()), raft::make_const_mdspan(reference_indices.view()), recall_value.view(), - raft::make_const_mdspan(distances), - raft::make_const_mdspan(reference_distances)); + raft::make_const_mdspan(distances.view()), + raft::make_const_mdspan(reference_distances.view())); res.sync_stream(); ``` @@ -340,4 +340,4 @@ The below example specifies the total number of bytes that RAFT can use for temp std::shared_ptr managed_resource; raft::device_resource res(managed_resource, std::make_optional(3 * 1024^3)); -``` \ No newline at end of file +``` From ff60d3a460229e1a6f57e9a3d32f47180bb53e78 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Thu, 30 Nov 2023 15:00:35 -0600 Subject: [PATCH 2/3] Update to fmt 10.1.1 and spdlog 1.12.0. (#1957) This PR updates to fmt 10.1.1 and spdlog 1.12. Depends on https://github.com/rapidsai/rmm/pull/1374. Authors: - Bradley Dice (https://github.com/bdice) Approvers: - Jake Awe (https://github.com/AyodeAwe) - Vyas Ramasubramani (https://github.com/vyasr) --- ci/build_python.sh | 12 ++++++------ conda/recipes/libraft/conda_build_config.yaml | 4 ++-- .../raft-ann-bench-cpu/conda_build_config.yaml | 4 ++-- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/ci/build_python.sh b/ci/build_python.sh index cf34776542..3e67edd5db 100755 --- a/ci/build_python.sh +++ b/ci/build_python.sh @@ -15,11 +15,11 @@ CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp) version=$(rapids-generate-version) git_commit=$(git rev-parse HEAD) -export RAPIDS_PACKAGE_VERSION=${version} +export RAPIDS_PACKAGE_VERSION=${version} echo "${version}" > VERSION package_dir="python" -for package_name in pylibraft raft-dask; do +for package_name in pylibraft raft-dask; do underscore_package_name=$(echo "${package_name}" | tr "-" "_") sed -i "/^__git_commit__/ s/= .*/= \"${git_commit}\"/g" "${package_dir}/${package_name}/${underscore_package_name}/_version.py" done @@ -39,10 +39,10 @@ rapids-conda-retry mambabuild \ # Build ann-bench for each cuda and python version rapids-conda-retry mambabuild \ ---no-test \ ---channel "${CPP_CHANNEL}" \ ---channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \ -conda/recipes/raft-ann-bench + --no-test \ + --channel "${CPP_CHANNEL}" \ + --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \ + conda/recipes/raft-ann-bench # Build ann-bench-cpu only in CUDA 11 jobs since it only depends on python # version diff --git a/conda/recipes/libraft/conda_build_config.yaml b/conda/recipes/libraft/conda_build_config.yaml index 25493a34fa..f1229281bb 100644 --- a/conda/recipes/libraft/conda_build_config.yaml +++ b/conda/recipes/libraft/conda_build_config.yaml @@ -73,7 +73,7 @@ cuda11_cuda_profiler_api_run_version: - ">=11.4.240,<12" spdlog_version: - - ">=1.11.0,<1.12" + - ">=1.12.0,<1.13" fmt_version: - - ">=9.1.0,<10" + - ">=10.1.1,<11" diff --git a/conda/recipes/raft-ann-bench-cpu/conda_build_config.yaml b/conda/recipes/raft-ann-bench-cpu/conda_build_config.yaml index fda3e4e53d..93a5532962 100644 --- a/conda/recipes/raft-ann-bench-cpu/conda_build_config.yaml +++ b/conda/recipes/raft-ann-bench-cpu/conda_build_config.yaml @@ -20,7 +20,7 @@ nlohmann_json_version: - ">=3.11.2" spdlog_version: - - ">=1.11.0,<1.12" + - ">=1.12.0,<1.13" fmt_version: - - ">=9.1.0,<10" + - ">=10.1.1,<11" From 4ba013951d5dfbbcbbe191896d88a43171ea057c Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Thu, 30 Nov 2023 15:02:41 -0800 Subject: [PATCH 3/3] Remove selection_faiss (#2027) Remove the selection_faiss instantiations. Since #1985, we haven't been using the faiss select_k code and these aren't necessary anymore. This should lead to a 70MB improvement in libraft.so binary size. This also removes the raft::spatial::select_k code in favour of matrix:: select_k - the spatial version was marked deprecated, and didn't switch between the best selection algorithms for the input size. Authors: - Ben Frederickson (https://github.com/benfred) Approvers: - Tamas Bela Feher (https://github.com/tfeher) URL: https://github.com/rapidsai/raft/pull/2027 --- cpp/CMakeLists.txt | 10 - cpp/bench/prims/matrix/select_k.cu | 3 - .../raft/matrix/detail/select_k-inl.cuh | 1 - .../neighbors/detail/selection_faiss-ext.cuh | 67 --- .../neighbors/detail/selection_faiss-inl.cuh | 163 ------ .../raft/neighbors/detail/selection_faiss.cuh | 24 - .../detail/selection_faiss_helpers.cuh | 31 -- .../raft/sparse/neighbors/detail/knn.cuh | 19 +- cpp/include/raft/spatial/knn/knn.cuh | 107 ---- .../raft_internal/matrix/select_k.cuh | 6 - .../detail/selection_faiss_00_generate.py | 79 --- .../detail/selection_faiss_int32_t_float.cu | 44 -- .../detail/selection_faiss_int64_t_double.cu | 44 -- .../detail/selection_faiss_int64_t_half.cu | 44 -- .../detail/selection_faiss_int_double.cu | 44 -- .../detail/selection_faiss_long_float.cu | 44 -- .../detail/selection_faiss_size_t_double.cu | 44 -- .../detail/selection_faiss_size_t_float.cu | 44 -- .../detail/selection_faiss_uint32_t_double.cu | 44 -- .../detail/selection_faiss_uint32_t_float.cu | 44 -- .../detail/selection_faiss_uint32_t_half.cu | 44 -- cpp/test/CMakeLists.txt | 6 - cpp/test/ext_headers/00_generate.py | 1 - .../raft_neighbors_detail_selection_faiss.cu | 27 - cpp/test/matrix/select_k.cuh | 2 - cpp/test/neighbors/fused_l2_knn.cu | 17 +- cpp/test/neighbors/selection.cu | 499 ------------------ cpp/test/neighbors/tiled_knn.cu | 18 +- 28 files changed, 25 insertions(+), 1495 deletions(-) delete mode 100644 cpp/include/raft/neighbors/detail/selection_faiss-ext.cuh delete mode 100644 cpp/include/raft/neighbors/detail/selection_faiss-inl.cuh delete mode 100644 cpp/include/raft/neighbors/detail/selection_faiss.cuh delete mode 100644 cpp/include/raft/neighbors/detail/selection_faiss_helpers.cuh delete mode 100644 cpp/src/neighbors/detail/selection_faiss_00_generate.py delete mode 100644 cpp/src/neighbors/detail/selection_faiss_int32_t_float.cu delete mode 100644 cpp/src/neighbors/detail/selection_faiss_int64_t_double.cu delete mode 100644 cpp/src/neighbors/detail/selection_faiss_int64_t_half.cu delete mode 100644 cpp/src/neighbors/detail/selection_faiss_int_double.cu delete mode 100644 cpp/src/neighbors/detail/selection_faiss_long_float.cu delete mode 100644 cpp/src/neighbors/detail/selection_faiss_size_t_double.cu delete mode 100644 cpp/src/neighbors/detail/selection_faiss_size_t_float.cu delete mode 100644 cpp/src/neighbors/detail/selection_faiss_uint32_t_double.cu delete mode 100644 cpp/src/neighbors/detail/selection_faiss_uint32_t_float.cu delete mode 100644 cpp/src/neighbors/detail/selection_faiss_uint32_t_half.cu delete mode 100644 cpp/test/ext_headers/raft_neighbors_detail_selection_faiss.cu delete mode 100644 cpp/test/neighbors/selection.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index acb77ec8c7..bccbc8c471 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -372,16 +372,6 @@ if(RAFT_COMPILE_LIBRARY) src/neighbors/detail/refine_host_float_float.cpp src/neighbors/detail/refine_host_int8_t_float.cpp src/neighbors/detail/refine_host_uint8_t_float.cpp - src/neighbors/detail/selection_faiss_int32_t_float.cu - src/neighbors/detail/selection_faiss_int_double.cu - src/neighbors/detail/selection_faiss_long_float.cu - src/neighbors/detail/selection_faiss_size_t_double.cu - src/neighbors/detail/selection_faiss_size_t_float.cu - src/neighbors/detail/selection_faiss_uint32_t_float.cu - src/neighbors/detail/selection_faiss_int64_t_double.cu - src/neighbors/detail/selection_faiss_int64_t_half.cu - src/neighbors/detail/selection_faiss_uint32_t_double.cu - src/neighbors/detail/selection_faiss_uint32_t_half.cu src/neighbors/ivf_flat_build_float_int64_t.cu src/neighbors/ivf_flat_build_int8_t_int64_t.cu src/neighbors/ivf_flat_build_uint8_t_int64_t.cu diff --git a/cpp/bench/prims/matrix/select_k.cu b/cpp/bench/prims/matrix/select_k.cu index d3994e59c5..324d3aef84 100644 --- a/cpp/bench/prims/matrix/select_k.cu +++ b/cpp/bench/prims/matrix/select_k.cu @@ -279,9 +279,6 @@ const static size_t MAX_MEMORY = 16 * 1024 * 1024 * 1024ULL; SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kWarpDistributed, input) \ SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kWarpDistributedShm, input) \ } \ - if (input.k <= raft::neighbors::detail::kFaissMaxK()) { \ - SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kFaissBlockSelect, input) \ - } \ } \ } diff --git a/cpp/include/raft/matrix/detail/select_k-inl.cuh b/cpp/include/raft/matrix/detail/select_k-inl.cuh index 20fe1963fc..9024975734 100644 --- a/cpp/include/raft/matrix/detail/select_k-inl.cuh +++ b/cpp/include/raft/matrix/detail/select_k-inl.cuh @@ -25,7 +25,6 @@ #include #include -#include #include #include #include diff --git a/cpp/include/raft/neighbors/detail/selection_faiss-ext.cuh b/cpp/include/raft/neighbors/detail/selection_faiss-ext.cuh deleted file mode 100644 index a6ed17e251..0000000000 --- a/cpp/include/raft/neighbors/detail/selection_faiss-ext.cuh +++ /dev/null @@ -1,67 +0,0 @@ -/* - * Copyright (c) 2020-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 // size_t -#include // uint32_t -#include // __half -#include // kFaissMaxK -#include // RAFT_EXPLICIT - -#if defined(RAFT_EXPLICIT_INSTANTIATE_ONLY) - -namespace raft::neighbors::detail { - -template -void select_k(const key_t* inK, - const payload_t* inV, - size_t n_rows, - size_t n_cols, - key_t* outK, - payload_t* outV, - bool select_min, - int k, - cudaStream_t stream) RAFT_EXPLICIT; -}; // namespace raft::neighbors::detail - -#endif // RAFT_EXPLICIT_INSTANTIATE_ONLY - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ - extern template void raft::neighbors::detail::select_k(const key_t* inK, \ - const payload_t* inV, \ - size_t n_rows, \ - size_t n_cols, \ - key_t* outK, \ - payload_t* outV, \ - bool select_min, \ - int k, \ - cudaStream_t stream) - -instantiate_raft_neighbors_detail_select_k(uint32_t, float); -instantiate_raft_neighbors_detail_select_k(int32_t, float); -instantiate_raft_neighbors_detail_select_k(long, float); -instantiate_raft_neighbors_detail_select_k(size_t, double); -// test/neighbors/selection.cu -instantiate_raft_neighbors_detail_select_k(int, double); -instantiate_raft_neighbors_detail_select_k(size_t, float); - -instantiate_raft_neighbors_detail_select_k(uint32_t, double); -instantiate_raft_neighbors_detail_select_k(int64_t, double); -instantiate_raft_neighbors_detail_select_k(uint32_t, __half); -instantiate_raft_neighbors_detail_select_k(int64_t, __half); - -#undef instantiate_raft_neighbors_detail_select_k diff --git a/cpp/include/raft/neighbors/detail/selection_faiss-inl.cuh b/cpp/include/raft/neighbors/detail/selection_faiss-inl.cuh deleted file mode 100644 index 35f322bbf6..0000000000 --- a/cpp/include/raft/neighbors/detail/selection_faiss-inl.cuh +++ /dev/null @@ -1,163 +0,0 @@ -/* - * Copyright (c) 2020-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 - -#include -#include // kFaissMaxK - -namespace raft::neighbors::detail { - -template -RAFT_KERNEL select_k_kernel(const key_t* inK, - const payload_t* inV, - size_t n_rows, - size_t n_cols, - key_t* outK, - payload_t* outV, - key_t initK, - payload_t initV, - int k) -{ - using align_warp = Pow2; - constexpr int kNumWarps = align_warp::div(tpb); - - __shared__ key_t smemK[kNumWarps * warp_q]; - __shared__ payload_t smemV[kNumWarps * warp_q]; - - faiss_select::BlockSelect, - warp_q, - thread_q, - tpb> - heap(initK, initV, smemK, smemV, k); - - // Grid is exactly sized to rows available - int row = blockIdx.x; - { - size_t i = size_t(threadIdx.x); - - inK += row * n_cols; - if (inV != nullptr) { inV += row * n_cols; } - - // Whole warps must participate in the selection - size_t limit = align_warp::roundDown(n_cols); - - for (; i < limit; i += tpb) { - heap.add(inK[i], (inV != nullptr) ? inV[i] : payload_t(i)); - } - - // Handle last remainder fraction of a warp of elements - if (i < n_cols) { heap.addThreadQ(inK[i], (inV != nullptr) ? inV[i] : payload_t(i)); } - } - - heap.reduce(); - - for (int i = threadIdx.x; i < k; i += tpb) { - outK[row * k + i] = smemK[i]; - outV[row * k + i] = smemV[i]; - } -} - -template -inline void select_k_impl(const key_t* inK, - const payload_t* inV, - size_t n_rows, - size_t n_cols, - key_t* outK, - payload_t* outV, - bool select_min, - int k, - cudaStream_t stream) -{ - auto grid = dim3(n_rows); - - constexpr int n_threads = (warp_q <= 1024) ? 128 : 64; - auto block = dim3(n_threads); - - auto kInit = select_min ? upper_bound() : lower_bound(); - auto vInit = -1; - if (select_min) { - select_k_kernel - <<>>(inK, inV, n_rows, n_cols, outK, outV, kInit, vInit, k); - } else { - select_k_kernel - <<>>(inK, inV, n_rows, n_cols, outK, outV, kInit, vInit, k); - } - RAFT_CUDA_TRY(cudaGetLastError()); -} - -/** - * @brief Select the k-nearest neighbors from dense - * distance and index matrices. - * - * @param[in] inK partitioned knn distance matrix - * @param[in] inV partitioned knn index matrix - * @param[in] n_rows number of rows in distance and index matrices - * @param[in] n_cols number of columns in distance and index matrices - * @param[out] outK merged knn distance matrix - * @param[out] outV merged knn index matrix - * @param[in] select_min whether to select the min or the max distances - * @param[in] k number of neighbors per partition (also number of merged neighbors) - * @param[in] stream CUDA stream to use - */ -template -inline void select_k(const key_t* inK, - const payload_t* inV, - size_t n_rows, - size_t n_cols, - key_t* outK, - payload_t* outV, - bool select_min, - int k, - cudaStream_t stream) -{ - constexpr int max_k = kFaissMaxK(); - if (k == 1) - select_k_impl( - inK, inV, n_rows, n_cols, outK, outV, select_min, k, stream); - else if (k <= 32) - select_k_impl( - inK, inV, n_rows, n_cols, outK, outV, select_min, k, stream); - else if (k <= 64) - select_k_impl( - inK, inV, n_rows, n_cols, outK, outV, select_min, k, stream); - else if (k <= 128) - select_k_impl( - inK, inV, n_rows, n_cols, outK, outV, select_min, k, stream); - else if (k <= 256) - select_k_impl( - inK, inV, n_rows, n_cols, outK, outV, select_min, k, stream); - else if (k <= 512) - select_k_impl( - inK, inV, n_rows, n_cols, outK, outV, select_min, k, stream); - else if (k <= 1024 && k <= max_k) - // note: have to use constexpr std::min here to avoid instantiating templates - // for parameters we don't support - select_k_impl( - inK, inV, n_rows, n_cols, outK, outV, select_min, k, stream); - else if (k <= 2048 && k <= max_k) - select_k_impl( - inK, inV, n_rows, n_cols, outK, outV, select_min, k, stream); - else - ASSERT(k <= max_k, "Current max k is %d (requested %d)", max_k, k); -} -}; // namespace raft::neighbors::detail diff --git a/cpp/include/raft/neighbors/detail/selection_faiss.cuh b/cpp/include/raft/neighbors/detail/selection_faiss.cuh deleted file mode 100644 index dd229b37e8..0000000000 --- a/cpp/include/raft/neighbors/detail/selection_faiss.cuh +++ /dev/null @@ -1,24 +0,0 @@ -/* - * Copyright (c) 2020-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 - -#ifndef RAFT_EXPLICIT_INSTANTIATE_ONLY -#include "selection_faiss-inl.cuh" -#endif - -#ifdef RAFT_COMPILED -#include "selection_faiss-ext.cuh" -#endif diff --git a/cpp/include/raft/neighbors/detail/selection_faiss_helpers.cuh b/cpp/include/raft/neighbors/detail/selection_faiss_helpers.cuh deleted file mode 100644 index c4b69f21ec..0000000000 --- a/cpp/include/raft/neighbors/detail/selection_faiss_helpers.cuh +++ /dev/null @@ -1,31 +0,0 @@ -/* - * Copyright (c) 2020-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 - -namespace raft::neighbors::detail { - -// This function is used in cpp/test/neighbors/select.cu. We want to make it -// available through both the selection_faiss-inl.cuh and -// selection_faiss-ext.cuh headers. -template -constexpr int kFaissMaxK() -{ - if (sizeof(key_t) >= 8) { return sizeof(payload_t) >= 8 ? 512 : 1024; } - return 2048; -} - -} // namespace raft::neighbors::detail diff --git a/cpp/include/raft/sparse/neighbors/detail/knn.cuh b/cpp/include/raft/sparse/neighbors/detail/knn.cuh index ff644c000e..f9d019931a 100644 --- a/cpp/include/raft/sparse/neighbors/detail/knn.cuh +++ b/cpp/include/raft/sparse/neighbors/detail/knn.cuh @@ -24,12 +24,12 @@ #include #include +#include #include #include #include #include #include -#include #include @@ -365,15 +365,14 @@ class sparse_knn_t { bool ascending = raft::distance::is_min_close(metric); // kernel to slice first (min) k cols and copy into batched merge buffer - raft::spatial::knn::select_k(batch_dists, - batch_indices, - batch_rows, - batch_cols, - out_dists, - out_indices, - ascending, - n_neighbors, - resource::get_cuda_stream(handle)); + raft::matrix::select_k( + handle, + make_device_matrix_view(batch_dists, batch_rows, batch_cols), + make_device_matrix_view(batch_indices, batch_rows, batch_cols), + make_device_matrix_view(out_dists, batch_rows, n_neighbors), + make_device_matrix_view(out_indices, batch_rows, n_neighbors), + ascending, + true); } void compute_distances(csr_batcher_t& idx_batcher, diff --git a/cpp/include/raft/spatial/knn/knn.cuh b/cpp/include/raft/spatial/knn/knn.cuh index 3c089b1d22..9ece7d3f1c 100644 --- a/cpp/include/raft/spatial/knn/knn.cuh +++ b/cpp/include/raft/spatial/knn/knn.cuh @@ -21,7 +21,6 @@ #include #include #include -#include namespace raft::spatial::knn { @@ -64,112 +63,6 @@ inline void knn_merge_parts(const value_t* in_keys, in_keys, in_values, out_keys, out_values, n_samples, n_parts, k, stream, translations); } -/** Choose an implementation for the select-top-k, */ -enum class SelectKAlgo { - /** Adapted from the faiss project. Result: sorted (not stable). */ - FAISS, - /** Incomplete series of radix sort passes, comparing 8 bits per pass. Result: unsorted. */ - RADIX_8_BITS, - /** Incomplete series of radix sort passes, comparing 11 bits per pass. Result: unsorted. */ - RADIX_11_BITS, - /** Filtering with a bitonic-sort-based priority queue. Result: sorted (not stable). */ - WARP_SORT -}; - -/** - * Select k smallest or largest key/values from each row in the input data. - * - * If you think of the input data `in_keys` as a row-major matrix with input_len columns and - * n_inputs rows, then this function selects k smallest/largest values in each row and fills - * in the row-major matrix `out_keys` of size (n_inputs, k). - * - * Note, depending on the selected algorithm, the values within rows of `out_keys` are not - * necessarily sorted. See the `SelectKAlgo` enumeration for more details. - * - * Note: This call is deprecated, please use `raft/matrix/select_k.cuh` - * - * @tparam idx_t - * the payload type (what is being selected together with the keys). - * @tparam value_t - * the type of the keys (what is being compared). - * - * @param[in] in_keys - * contiguous device array of inputs of size (input_len * n_inputs); - * these are compared and selected. - * @param[in] in_values - * contiguous device array of inputs of size (input_len * n_inputs); - * typically, these are indices of the corresponding in_keys. - * You can pass `NULL` as an argument here; this would imply `in_values` is a homogeneous array - * of indices from `0` to `input_len - 1` for every input and reduce the usage of memory - * bandwidth. - * @param[in] n_inputs - * number of input rows, i.e. the batch size. - * @param[in] input_len - * length of a single input array (row); also sometimes referred as n_cols. - * Invariant: input_len >= k. - * @param[out] out_keys - * contiguous device array of outputs of size (k * n_inputs); - * the k smallest/largest values from each row of the `in_keys`. - * @param[out] out_values - * contiguous device array of outputs of size (k * n_inputs); - * the payload selected together with `out_keys`. - * @param[in] select_min - * whether to select k smallest (true) or largest (false) keys. - * @param[in] k - * the number of outputs to select in each input row. - * @param[in] stream - * @param[in] algo - * the implementation of the algorithm - */ -template -[[deprecated("Use function `select_k` from `raft/matrix/select_k.cuh`")]] inline void select_k( - const value_t* in_keys, - const idx_t* in_values, - size_t n_inputs, - size_t input_len, - value_t* out_keys, - idx_t* out_values, - bool select_min, - int k, - cudaStream_t stream, - SelectKAlgo algo = SelectKAlgo::FAISS) -{ - common::nvtx::range fun_scope("select-%s-%d (%zu, %zu) algo-%d", - select_min ? "min" : "max", - k, - n_inputs, - input_len, - int(algo)); - ASSERT(size_t(input_len) >= size_t(k), - "Size of the input (input_len = %zu) must be not smaller than the selection (k = %zu).", - size_t(input_len), - size_t(k)); - - switch (algo) { - case SelectKAlgo::FAISS: - neighbors::detail::select_k( - in_keys, in_values, n_inputs, input_len, out_keys, out_values, select_min, k, stream); - break; - - case SelectKAlgo::RADIX_8_BITS: - matrix::detail::select::radix::select_k( - in_keys, in_values, n_inputs, input_len, k, out_keys, out_values, select_min, true, stream); - break; - - case SelectKAlgo::RADIX_11_BITS: - matrix::detail::select::radix::select_k( - in_keys, in_values, n_inputs, input_len, k, out_keys, out_values, select_min, true, stream); - break; - - case SelectKAlgo::WARP_SORT: - matrix::detail::select::warpsort::select_k( - in_keys, in_values, n_inputs, input_len, k, out_keys, out_values, select_min, stream); - break; - - default: ASSERT(false, "Unknown algorithm (id = %d)", int(algo)); - } -} - /** * @brief Flat C++ API function to perform a brute force knn on * a series of input arrays and combine the results into a single diff --git a/cpp/internal/raft_internal/matrix/select_k.cuh b/cpp/internal/raft_internal/matrix/select_k.cuh index 1d15c5fc03..93095ff82e 100644 --- a/cpp/internal/raft_internal/matrix/select_k.cuh +++ b/cpp/internal/raft_internal/matrix/select_k.cuh @@ -21,7 +21,6 @@ #include #include #include -#include namespace raft::matrix::select { @@ -59,7 +58,6 @@ enum class Algo { kWarpFiltered, kWarpDistributed, kWarpDistributedShm, - kFaissBlockSelect }; inline auto operator<<(std::ostream& os, const Algo& algo) -> std::ostream& @@ -74,7 +72,6 @@ inline auto operator<<(std::ostream& os, const Algo& algo) -> std::ostream& case Algo::kWarpFiltered: return os << "kWarpFiltered"; case Algo::kWarpDistributed: return os << "kWarpDistributed"; case Algo::kWarpDistributedShm: return os << "kWarpDistributedShm"; - case Algo::kFaissBlockSelect: return os << "kFaissBlockSelect"; default: return os << "unknown enum value"; } } @@ -167,9 +164,6 @@ void select_k_impl(const resources& handle, return detail::select::warpsort:: select_k_impl( in, in_idx, batch_size, len, k, out, out_idx, select_min, stream); - case Algo::kFaissBlockSelect: - return neighbors::detail::select_k( - in, in_idx, batch_size, len, out, out_idx, select_min, k, stream); } } } // namespace raft::matrix::select diff --git a/cpp/src/neighbors/detail/selection_faiss_00_generate.py b/cpp/src/neighbors/detail/selection_faiss_00_generate.py deleted file mode 100644 index 386dd18e0c..0000000000 --- a/cpp/src/neighbors/detail/selection_faiss_00_generate.py +++ /dev/null @@ -1,79 +0,0 @@ -# Copyright (c) 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. - -header = """ -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by selection_faiss_00_generate.py - * - * Make changes there and run in this directory: - * - * > python selection_faiss_00_generate.py - * - */ - -#include // size_t -#include // uint32_t -#include - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \\ - template void raft::neighbors::detail::select_k(const key_t* inK, \\ - const payload_t* inV, \\ - size_t n_rows, \\ - size_t n_cols, \\ - key_t* outK, \\ - payload_t* outV, \\ - bool select_min, \\ - int k, \\ - cudaStream_t stream) - -""" - -types = dict( - uint32_t_float=("uint32_t", "float"), - uint32_t_double=("uint32_t", "double"), - uint32_t_half=("uint32_t", "half"), - int64_t_double=("int64_t", "double"), - int64_t_half=("int64_t", "half"), - int32_t_float=("int32_t", "float"), - long_float=("long", "float"), - size_t_double=("size_t", "double"), - int_double=("int", "double"), - size_t_float=("size_t", "float"), -) - -for type_path, (payload_t, key_t) in types.items(): - path = f"selection_faiss_{type_path}.cu" - with open(path, "w") as f: - f.write(header) - f.write(f"instantiate_raft_neighbors_detail_select_k({payload_t}, {key_t});\n\n") - f.write(f"#undef instantiate_raft_neighbors_detail_select_k\n") - - # for pasting into CMakeLists.txt - print(f"src/neighbors/detail/{path}") diff --git a/cpp/src/neighbors/detail/selection_faiss_int32_t_float.cu b/cpp/src/neighbors/detail/selection_faiss_int32_t_float.cu deleted file mode 100644 index 1f1ece05ae..0000000000 --- a/cpp/src/neighbors/detail/selection_faiss_int32_t_float.cu +++ /dev/null @@ -1,44 +0,0 @@ - -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by selection_faiss_00_generate.py - * - * Make changes there and run in this directory: - * - * > python selection_faiss_00_generate.py - * - */ - -#include // size_t -#include // uint32_t -#include - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ - template void raft::neighbors::detail::select_k(const key_t* inK, \ - const payload_t* inV, \ - size_t n_rows, \ - size_t n_cols, \ - key_t* outK, \ - payload_t* outV, \ - bool select_min, \ - int k, \ - cudaStream_t stream) - -instantiate_raft_neighbors_detail_select_k(int32_t, float); - -#undef instantiate_raft_neighbors_detail_select_k diff --git a/cpp/src/neighbors/detail/selection_faiss_int64_t_double.cu b/cpp/src/neighbors/detail/selection_faiss_int64_t_double.cu deleted file mode 100644 index f824fdd479..0000000000 --- a/cpp/src/neighbors/detail/selection_faiss_int64_t_double.cu +++ /dev/null @@ -1,44 +0,0 @@ - -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by selection_faiss_00_generate.py - * - * Make changes there and run in this directory: - * - * > python selection_faiss_00_generate.py - * - */ - -#include // size_t -#include // uint32_t -#include - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ - template void raft::neighbors::detail::select_k(const key_t* inK, \ - const payload_t* inV, \ - size_t n_rows, \ - size_t n_cols, \ - key_t* outK, \ - payload_t* outV, \ - bool select_min, \ - int k, \ - cudaStream_t stream) - -instantiate_raft_neighbors_detail_select_k(int64_t, double); - -#undef instantiate_raft_neighbors_detail_select_k diff --git a/cpp/src/neighbors/detail/selection_faiss_int64_t_half.cu b/cpp/src/neighbors/detail/selection_faiss_int64_t_half.cu deleted file mode 100644 index 34ca525c64..0000000000 --- a/cpp/src/neighbors/detail/selection_faiss_int64_t_half.cu +++ /dev/null @@ -1,44 +0,0 @@ - -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by selection_faiss_00_generate.py - * - * Make changes there and run in this directory: - * - * > python selection_faiss_00_generate.py - * - */ - -#include // size_t -#include // uint32_t -#include - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ - template void raft::neighbors::detail::select_k(const key_t* inK, \ - const payload_t* inV, \ - size_t n_rows, \ - size_t n_cols, \ - key_t* outK, \ - payload_t* outV, \ - bool select_min, \ - int k, \ - cudaStream_t stream) - -instantiate_raft_neighbors_detail_select_k(int64_t, half); - -#undef instantiate_raft_neighbors_detail_select_k diff --git a/cpp/src/neighbors/detail/selection_faiss_int_double.cu b/cpp/src/neighbors/detail/selection_faiss_int_double.cu deleted file mode 100644 index 7e832410c4..0000000000 --- a/cpp/src/neighbors/detail/selection_faiss_int_double.cu +++ /dev/null @@ -1,44 +0,0 @@ - -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by selection_faiss_00_generate.py - * - * Make changes there and run in this directory: - * - * > python selection_faiss_00_generate.py - * - */ - -#include // size_t -#include // uint32_t -#include - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ - template void raft::neighbors::detail::select_k(const key_t* inK, \ - const payload_t* inV, \ - size_t n_rows, \ - size_t n_cols, \ - key_t* outK, \ - payload_t* outV, \ - bool select_min, \ - int k, \ - cudaStream_t stream) - -instantiate_raft_neighbors_detail_select_k(int, double); - -#undef instantiate_raft_neighbors_detail_select_k diff --git a/cpp/src/neighbors/detail/selection_faiss_long_float.cu b/cpp/src/neighbors/detail/selection_faiss_long_float.cu deleted file mode 100644 index 441d54fa30..0000000000 --- a/cpp/src/neighbors/detail/selection_faiss_long_float.cu +++ /dev/null @@ -1,44 +0,0 @@ - -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by selection_faiss_00_generate.py - * - * Make changes there and run in this directory: - * - * > python selection_faiss_00_generate.py - * - */ - -#include // size_t -#include // uint32_t -#include - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ - template void raft::neighbors::detail::select_k(const key_t* inK, \ - const payload_t* inV, \ - size_t n_rows, \ - size_t n_cols, \ - key_t* outK, \ - payload_t* outV, \ - bool select_min, \ - int k, \ - cudaStream_t stream) - -instantiate_raft_neighbors_detail_select_k(long, float); - -#undef instantiate_raft_neighbors_detail_select_k diff --git a/cpp/src/neighbors/detail/selection_faiss_size_t_double.cu b/cpp/src/neighbors/detail/selection_faiss_size_t_double.cu deleted file mode 100644 index ca310e7697..0000000000 --- a/cpp/src/neighbors/detail/selection_faiss_size_t_double.cu +++ /dev/null @@ -1,44 +0,0 @@ - -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by selection_faiss_00_generate.py - * - * Make changes there and run in this directory: - * - * > python selection_faiss_00_generate.py - * - */ - -#include // size_t -#include // uint32_t -#include - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ - template void raft::neighbors::detail::select_k(const key_t* inK, \ - const payload_t* inV, \ - size_t n_rows, \ - size_t n_cols, \ - key_t* outK, \ - payload_t* outV, \ - bool select_min, \ - int k, \ - cudaStream_t stream) - -instantiate_raft_neighbors_detail_select_k(size_t, double); - -#undef instantiate_raft_neighbors_detail_select_k diff --git a/cpp/src/neighbors/detail/selection_faiss_size_t_float.cu b/cpp/src/neighbors/detail/selection_faiss_size_t_float.cu deleted file mode 100644 index a830e6ecac..0000000000 --- a/cpp/src/neighbors/detail/selection_faiss_size_t_float.cu +++ /dev/null @@ -1,44 +0,0 @@ - -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by selection_faiss_00_generate.py - * - * Make changes there and run in this directory: - * - * > python selection_faiss_00_generate.py - * - */ - -#include // size_t -#include // uint32_t -#include - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ - template void raft::neighbors::detail::select_k(const key_t* inK, \ - const payload_t* inV, \ - size_t n_rows, \ - size_t n_cols, \ - key_t* outK, \ - payload_t* outV, \ - bool select_min, \ - int k, \ - cudaStream_t stream) - -instantiate_raft_neighbors_detail_select_k(size_t, float); - -#undef instantiate_raft_neighbors_detail_select_k diff --git a/cpp/src/neighbors/detail/selection_faiss_uint32_t_double.cu b/cpp/src/neighbors/detail/selection_faiss_uint32_t_double.cu deleted file mode 100644 index e39edbb031..0000000000 --- a/cpp/src/neighbors/detail/selection_faiss_uint32_t_double.cu +++ /dev/null @@ -1,44 +0,0 @@ - -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by selection_faiss_00_generate.py - * - * Make changes there and run in this directory: - * - * > python selection_faiss_00_generate.py - * - */ - -#include // size_t -#include // uint32_t -#include - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ - template void raft::neighbors::detail::select_k(const key_t* inK, \ - const payload_t* inV, \ - size_t n_rows, \ - size_t n_cols, \ - key_t* outK, \ - payload_t* outV, \ - bool select_min, \ - int k, \ - cudaStream_t stream) - -instantiate_raft_neighbors_detail_select_k(uint32_t, double); - -#undef instantiate_raft_neighbors_detail_select_k diff --git a/cpp/src/neighbors/detail/selection_faiss_uint32_t_float.cu b/cpp/src/neighbors/detail/selection_faiss_uint32_t_float.cu deleted file mode 100644 index 2fecaa5cf1..0000000000 --- a/cpp/src/neighbors/detail/selection_faiss_uint32_t_float.cu +++ /dev/null @@ -1,44 +0,0 @@ - -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by selection_faiss_00_generate.py - * - * Make changes there and run in this directory: - * - * > python selection_faiss_00_generate.py - * - */ - -#include // size_t -#include // uint32_t -#include - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ - template void raft::neighbors::detail::select_k(const key_t* inK, \ - const payload_t* inV, \ - size_t n_rows, \ - size_t n_cols, \ - key_t* outK, \ - payload_t* outV, \ - bool select_min, \ - int k, \ - cudaStream_t stream) - -instantiate_raft_neighbors_detail_select_k(uint32_t, float); - -#undef instantiate_raft_neighbors_detail_select_k diff --git a/cpp/src/neighbors/detail/selection_faiss_uint32_t_half.cu b/cpp/src/neighbors/detail/selection_faiss_uint32_t_half.cu deleted file mode 100644 index 549509f6da..0000000000 --- a/cpp/src/neighbors/detail/selection_faiss_uint32_t_half.cu +++ /dev/null @@ -1,44 +0,0 @@ - -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by selection_faiss_00_generate.py - * - * Make changes there and run in this directory: - * - * > python selection_faiss_00_generate.py - * - */ - -#include // size_t -#include // uint32_t -#include - -#define instantiate_raft_neighbors_detail_select_k(payload_t, key_t) \ - template void raft::neighbors::detail::select_k(const key_t* inK, \ - const payload_t* inV, \ - size_t n_rows, \ - size_t n_cols, \ - key_t* outK, \ - payload_t* outV, \ - bool select_min, \ - int k, \ - cudaStream_t stream) - -instantiate_raft_neighbors_detail_select_k(uint32_t, half); - -#undef instantiate_raft_neighbors_detail_select_k diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt index 6c03da8d7f..847dec8568 100644 --- a/cpp/test/CMakeLists.txt +++ b/cpp/test/CMakeLists.txt @@ -190,7 +190,6 @@ if(BUILD_TESTS) test/ext_headers/raft_core_logger.cpp test/ext_headers/raft_neighbors_refine.cu test/ext_headers/raft_neighbors_detail_ivf_flat_search.cu - test/ext_headers/raft_neighbors_detail_selection_faiss.cu test/ext_headers/raft_linalg_detail_coalesced_reduction.cu test/ext_headers/raft_spatial_knn_detail_ball_cover_registers.cu test/ext_headers/raft_neighbors_detail_ivf_flat_interleaved_scan.cu @@ -412,11 +411,6 @@ if(BUILD_TESTS) 100 ) - ConfigureTest( - NAME NEIGHBORS_SELECTION_TEST PATH test/neighbors/selection.cu LIB EXPLICIT_INSTANTIATE_ONLY - GPUS 1 PERCENT 50 - ) - ConfigureTest( NAME STATS_TEST diff --git a/cpp/test/ext_headers/00_generate.py b/cpp/test/ext_headers/00_generate.py index 15f90e1cc5..682cadbe89 100644 --- a/cpp/test/ext_headers/00_generate.py +++ b/cpp/test/ext_headers/00_generate.py @@ -54,7 +54,6 @@ "raft/core/logger-ext.hpp", "raft/neighbors/refine-ext.cuh", "raft/neighbors/detail/ivf_flat_search-ext.cuh", - "raft/neighbors/detail/selection_faiss-ext.cuh", "raft/linalg/detail/coalesced_reduction-ext.cuh", "raft/spatial/knn/detail/ball_cover/registers-ext.cuh", "raft/neighbors/detail/ivf_flat_interleaved_scan-ext.cuh", diff --git a/cpp/test/ext_headers/raft_neighbors_detail_selection_faiss.cu b/cpp/test/ext_headers/raft_neighbors_detail_selection_faiss.cu deleted file mode 100644 index f8bd21e86f..0000000000 --- a/cpp/test/ext_headers/raft_neighbors_detail_selection_faiss.cu +++ /dev/null @@ -1,27 +0,0 @@ - -/* - * Copyright (c) 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. - */ - -/* - * NOTE: this file is generated by 00_generate.py - * - * Make changes there and run in this directory: - * - * > python 00_generate.py - * - */ - -#include diff --git a/cpp/test/matrix/select_k.cuh b/cpp/test/matrix/select_k.cuh index e94a6d029e..fdea982d6c 100644 --- a/cpp/test/matrix/select_k.cuh +++ b/cpp/test/matrix/select_k.cuh @@ -267,8 +267,6 @@ struct SelectK // NOLINT case select::Algo::kWarpFiltered: case select::Algo::kWarpDistributed: case select::Algo::kWarpDistributedShm: return ix == 0; - // FAISS version returns a special invalid value: - case select::Algo::kFaissBlockSelect: return ix == std::numeric_limits::max(); // Do not forgive by default default: return false; } diff --git a/cpp/test/neighbors/fused_l2_knn.cu b/cpp/test/neighbors/fused_l2_knn.cu index fd89dc0fc7..80a01ce568 100644 --- a/cpp/test/neighbors/fused_l2_knn.cu +++ b/cpp/test/neighbors/fused_l2_knn.cu @@ -83,15 +83,14 @@ class FusedL2KNNTest : public ::testing::TestWithParam { raft::make_device_matrix_view(temp_distances.data(), num_queries, num_db_vecs), metric); - spatial::knn::select_k(temp_distances.data(), - nullptr, - num_queries, - num_db_vecs, - ref_distances_.data(), - ref_indices_.data(), - true, - k_, - stream_); + matrix::select_k( + handle_, + make_device_matrix_view(temp_distances.data(), num_queries, num_db_vecs), + std::nullopt, + make_device_matrix_view(ref_distances_.data(), num_queries, k_), + make_device_matrix_view(ref_indices_.data(), num_queries, k_), + true, + true); auto index_view = raft::make_device_matrix_view(database.data(), num_db_vecs, dim); diff --git a/cpp/test/neighbors/selection.cu b/cpp/test/neighbors/selection.cu deleted file mode 100644 index 6030e2a1a6..0000000000 --- a/cpp/test/neighbors/selection.cu +++ /dev/null @@ -1,499 +0,0 @@ -/* - * Copyright (c) 2018-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. - */ - -#include -#include -#include -#include -#include -#include // kFaissMax -#include -#include - -#include "../test_utils.cuh" - -#include -#include - -namespace raft::spatial::selection { - -using namespace raft; -using namespace raft::sparse; - -struct SelectTestSpec { - int n_inputs; - int input_len; - int k; - int select_min; - bool use_index_input = true; -}; - -std::ostream& operator<<(std::ostream& os, const SelectTestSpec& ss) -{ - os << "spec{size: " << ss.input_len << "*" << ss.n_inputs << ", k: " << ss.k; - os << (ss.select_min ? "; min}" : "; max}"); - return os; -} - -template -auto gen_simple_ids(int n_inputs, int input_len, const raft::resources& handle) -> std::vector -{ - std::vector out(n_inputs * input_len); - auto s = resource::get_cuda_stream(handle); - rmm::device_uvector out_d(out.size(), s); - iota_fill(out_d.data(), IdxT(n_inputs), IdxT(input_len), s); - update_host(out.data(), out_d.data(), out.size(), s); - s.synchronize(); - return out; -} - -template -struct SelectInOutSimple { - public: - bool not_supported = false; - - SelectInOutSimple(std::shared_ptr handle, - const SelectTestSpec& spec, - const std::vector& in_dists, - const std::vector& out_dists, - const std::vector& out_ids) - : in_dists_(in_dists), - in_ids_(gen_simple_ids(spec.n_inputs, spec.input_len, *handle.get())), - out_dists_(out_dists), - out_ids_(out_ids), - handle_(handle) - { - } - - auto get_in_dists() -> std::vector& { return in_dists_; } - auto get_in_ids() -> std::vector& { return in_ids_; } - auto get_out_dists() -> std::vector& { return out_dists_; } - auto get_out_ids() -> std::vector& { return out_ids_; } - - private: - std::shared_ptr handle_; - std::vector in_dists_; - std::vector in_ids_; - std::vector out_dists_; - std::vector out_ids_; -}; - -template -struct SelectInOutComputed { - public: - bool not_supported = false; - - SelectInOutComputed(std::shared_ptr handle, - const SelectTestSpec& spec, - knn::SelectKAlgo algo, - const std::vector& in_dists, - const std::optional>& in_ids = std::nullopt) - : handle_(handle), - in_dists_(in_dists), - in_ids_(in_ids.value_or(gen_simple_ids(spec.n_inputs, spec.input_len, *handle.get()))), - out_dists_(spec.n_inputs * spec.k), - out_ids_(spec.n_inputs * spec.k) - - { - // check if the size is supported by the algorithm - switch (algo) { - case knn::SelectKAlgo::WARP_SORT: - if (spec.k > raft::matrix::detail::select::warpsort::kMaxCapacity) { - not_supported = true; - return; - } - break; - case knn::SelectKAlgo::FAISS: - if (spec.k > raft::neighbors::detail::kFaissMaxK()) { - not_supported = true; - return; - } - break; - default: break; - } - - auto stream = resource::get_cuda_stream(*handle_); - - rmm::device_uvector in_dists_d(in_dists_.size(), stream); - rmm::device_uvector in_ids_d(in_ids_.size(), stream); - rmm::device_uvector out_dists_d(out_dists_.size(), stream); - rmm::device_uvector out_ids_d(out_ids_.size(), stream); - - update_device(in_dists_d.data(), in_dists_.data(), in_dists_.size(), stream); - update_device(in_ids_d.data(), in_ids_.data(), in_ids_.size(), stream); - - raft::spatial::knn::select_k(in_dists_d.data(), - spec.use_index_input ? in_ids_d.data() : nullptr, - spec.n_inputs, - spec.input_len, - out_dists_d.data(), - out_ids_d.data(), - spec.select_min, - spec.k, - stream, - algo); - - update_host(out_dists_.data(), out_dists_d.data(), out_dists_.size(), stream); - update_host(out_ids_.data(), out_ids_d.data(), out_ids_.size(), stream); - - interruptible::synchronize(stream); - - auto p = topk_sort_permutation(out_dists_, out_ids_, spec.k, spec.select_min); - apply_permutation(out_dists_, p); - apply_permutation(out_ids_, p); - } - - auto get_in_dists() -> std::vector& { return in_dists_; } - auto get_in_ids() -> std::vector& { return in_ids_; } - auto get_out_dists() -> std::vector& { return out_dists_; } - auto get_out_ids() -> std::vector& { return out_ids_; } - - private: - std::shared_ptr handle_; - std::vector in_dists_; - std::vector in_ids_; - std::vector out_dists_; - std::vector out_ids_; - - auto topk_sort_permutation(const std::vector& vec, - const std::vector& inds, - int k, - bool select_min) -> std::vector - { - std::vector p(vec.size()); - std::iota(p.begin(), p.end(), 0); - if (select_min) { - std::sort(p.begin(), p.end(), [&vec, &inds, k](IdxT i, IdxT j) { - const IdxT ik = i / k; - const IdxT jk = j / k; - if (ik == jk) { - if (vec[i] == vec[j]) { return inds[i] < inds[j]; } - return vec[i] < vec[j]; - } - return ik < jk; - }); - } else { - std::sort(p.begin(), p.end(), [&vec, &inds, k](IdxT i, IdxT j) { - const IdxT ik = i / k; - const IdxT jk = j / k; - if (ik == jk) { - if (vec[i] == vec[j]) { return inds[i] < inds[j]; } - return vec[i] > vec[j]; - } - return ik < jk; - }); - } - return p; - } - - template - void apply_permutation(std::vector& vec, const std::vector& p) - { - for (auto i = IdxT(vec.size()) - 1; i > 0; i--) { - auto j = p[i]; - while (j > i) - j = p[j]; - std::swap(vec[j], vec[i]); - } - } -}; - -template -using Params = - std::tuple>; - -template typename ParamsReader> -class SelectionTest : public testing::TestWithParam::ParamsIn> { - protected: - std::shared_ptr handle_; - const SelectTestSpec spec; - const knn::SelectKAlgo algo; - - typename ParamsReader::InOut ref; - SelectInOutComputed res; - - public: - explicit SelectionTest(Params::InOut> ps) - : handle_(std::get<3>(ps)), - spec(std::get<0>(ps)), - algo(std::get<1>(ps)), - ref(std::get<2>(ps)), - res(handle_, spec, algo, ref.get_in_dists(), ref.get_in_ids()) - { - } - - explicit SelectionTest(typename ParamsReader::ParamsIn ps) - : SelectionTest(ParamsReader::read(ps)) - { - } - - SelectionTest() - : SelectionTest(testing::TestWithParam::ParamsIn>::GetParam()) - { - } - - void run() - { - if (ref.not_supported || res.not_supported) { GTEST_SKIP(); } - - ASSERT_TRUE(hostVecMatch(ref.get_out_dists(), res.get_out_dists(), Compare())); - // If the dists (keys) are the same, different corresponding ids may end up in the selection due - // to non-deterministic nature of some implementations. - auto& in_ids = ref.get_in_ids(); - auto& in_dists = ref.get_in_dists(); - - auto compare_ids = [&in_ids, &in_dists](const IdxT& i, const IdxT& j) { - if (i == j) return true; - auto ix_i = size_t(std::find(in_ids.begin(), in_ids.end(), i) - in_ids.begin()); - auto ix_j = size_t(std::find(in_ids.begin(), in_ids.end(), j) - in_ids.begin()); - if (ix_i >= in_ids.size() || ix_j >= in_ids.size()) return false; - auto dist_i = in_dists[ix_i]; - auto dist_j = in_dists[ix_j]; - if (dist_i == dist_j) return true; - std::cout << "ERROR: ref[" << ix_i << "] = " << dist_i << " != " - << "res[" << ix_j << "] = " << dist_j << std::endl; - return false; - }; - ASSERT_TRUE(hostVecMatch(ref.get_out_ids(), res.get_out_ids(), compare_ids)); - } -}; - -template -struct params_simple { - using InOut = SelectInOutSimple; - using Inputs = - std::tuple, std::vector, std::vector>; - using Handle = std::shared_ptr; - using ParamsIn = std::tuple; - - static auto read(ParamsIn ps) -> Params - { - auto ins = std::get<0>(ps); - auto algo = std::get<1>(ps); - auto handle = std::get<2>(ps); - return std::make_tuple( - std::get<0>(ins), - algo, - SelectInOutSimple( - handle, std::get<0>(ins), std::get<1>(ins), std::get<2>(ins), std::get<3>(ins)), - handle); - } -}; - -auto inputs_simple_f = testing::Values( - params_simple::Inputs( - {5, 5, 5, true, true}, - {5.0, 4.0, 3.0, 2.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, 2.0, 3.0, 5.0, - 1.0, 4.0, 5.0, 3.0, 2.0, 4.0, 1.0, 1.0, 3.0, 2.0, 5.0, 4.0}, - {1.0, 2.0, 3.0, 4.0, 5.0, 1.0, 2.0, 3.0, 4.0, 5.0, 1.0, 2.0, 3.0, - 4.0, 5.0, 1.0, 2.0, 3.0, 4.0, 5.0, 1.0, 2.0, 3.0, 4.0, 5.0}, - {4, 3, 2, 1, 0, 0, 1, 2, 3, 4, 3, 0, 1, 4, 2, 4, 2, 1, 3, 0, 0, 2, 1, 4, 3}), - params_simple::Inputs( - {5, 5, 3, true, true}, - {5.0, 4.0, 3.0, 2.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, 2.0, 3.0, 5.0, - 1.0, 4.0, 5.0, 3.0, 2.0, 4.0, 1.0, 1.0, 3.0, 2.0, 5.0, 4.0}, - {1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0}, - {4, 3, 2, 0, 1, 2, 3, 0, 1, 4, 2, 1, 0, 2, 1}), - params_simple::Inputs( - {5, 5, 5, true, false}, - {5.0, 4.0, 3.0, 2.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, 2.0, 3.0, 5.0, - 1.0, 4.0, 5.0, 3.0, 2.0, 4.0, 1.0, 1.0, 3.0, 2.0, 5.0, 4.0}, - {1.0, 2.0, 3.0, 4.0, 5.0, 1.0, 2.0, 3.0, 4.0, 5.0, 1.0, 2.0, 3.0, - 4.0, 5.0, 1.0, 2.0, 3.0, 4.0, 5.0, 1.0, 2.0, 3.0, 4.0, 5.0}, - {4, 3, 2, 1, 0, 0, 1, 2, 3, 4, 3, 0, 1, 4, 2, 4, 2, 1, 3, 0, 0, 2, 1, 4, 3}), - params_simple::Inputs( - {5, 5, 3, true, false}, - {5.0, 4.0, 3.0, 2.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, 2.0, 3.0, 5.0, - 1.0, 4.0, 5.0, 3.0, 2.0, 4.0, 1.0, 1.0, 3.0, 2.0, 5.0, 4.0}, - {1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 1.0, 2.0, 3.0}, - {4, 3, 2, 0, 1, 2, 3, 0, 1, 4, 2, 1, 0, 2, 1}), - params_simple::Inputs( - {5, 7, 3, true, true}, - {5.0, 4.0, 3.0, 2.0, 1.3, 7.5, 19.0, 9.0, 2.0, 3.0, 3.0, 5.0, 6.0, 4.0, 2.0, 3.0, 5.0, 1.0, - 4.0, 1.0, 1.0, 5.0, 7.0, 2.5, 4.0, 7.0, 8.0, 8.0, 1.0, 3.0, 2.0, 5.0, 4.0, 1.1, 1.2}, - {1.3, 2.0, 3.0, 2.0, 3.0, 3.0, 1.0, 1.0, 1.0, 2.5, 4.0, 5.0, 1.0, 1.1, 1.2}, - {4, 3, 2, 1, 2, 3, 3, 5, 6, 2, 3, 0, 0, 5, 6}), - params_simple::Inputs( - {1, 7, 3, true, true}, {2.0, 3.0, 5.0, 1.0, 4.0, 1.0, 1.0}, {1.0, 1.0, 1.0}, {3, 5, 6}), - params_simple::Inputs( - {1, 7, 3, false, false}, {2.0, 3.0, 5.0, 1.0, 4.0, 1.0, 1.0}, {5.0, 4.0, 3.0}, {2, 4, 1}), - params_simple::Inputs( - {1, 7, 3, false, true}, {2.0, 3.0, 5.0, 9.0, 4.0, 9.0, 9.0}, {9.0, 9.0, 9.0}, {3, 5, 6}), - params_simple::Inputs( - {1, 130, 5, false, true}, - {19, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, - 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, - 0, 1, 0, 1, 0, 1, 0, 1, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, - 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 3, 4, - 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 4, 4, 2, 3, 2, 3, 2, 3, 2, 3, 2, 20}, - {20, 19, 18, 17, 16}, - {129, 0, 117, 116, 115}), - params_simple::Inputs( - {1, 130, 15, false, true}, - {19, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, - 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, - 0, 1, 0, 1, 0, 1, 0, 1, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, - 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 3, 4, - 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 4, 4, 2, 3, 2, 3, 2, 3, 2, 3, 2, 20}, - {20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6}, - {129, 0, 117, 116, 115, 114, 113, 112, 111, 110, 109, 108, 107, 106, 105})); - -typedef SelectionTest SimpleFloatInt; -TEST_P(SimpleFloatInt, Run) { run(); } -INSTANTIATE_TEST_CASE_P(SelectionTest, - SimpleFloatInt, - testing::Combine(inputs_simple_f, - testing::Values(knn::SelectKAlgo::FAISS, - knn::SelectKAlgo::RADIX_8_BITS, - knn::SelectKAlgo::RADIX_11_BITS, - knn::SelectKAlgo::WARP_SORT), - testing::Values(std::make_shared()))); - -template -struct with_ref { - template - struct params_random { - using InOut = SelectInOutComputed; - using Handle = std::shared_ptr; - using ParamsIn = std::tuple; - - static auto read(ParamsIn ps) -> Params - { - auto spec = std::get<0>(ps); - auto algo = std::get<1>(ps); - auto handle = std::get<2>(ps); - - std::vector dists(spec.input_len * spec.n_inputs); - - { - auto s = resource::get_cuda_stream(*handle); - rmm::device_uvector dists_d(spec.input_len * spec.n_inputs, s); - raft::random::RngState r(42); - normal(*(handle.get()), r, dists_d.data(), dists_d.size(), KeyT(10.0), KeyT(100.0)); - update_host(dists.data(), dists_d.data(), dists_d.size(), s); - s.synchronize(); - } - - return std::make_tuple( - spec, algo, SelectInOutComputed(handle, spec, RefAlgo, dists), handle); - } - }; -}; - -auto inputs_random_longlist = testing::Values(SelectTestSpec{1, 130, 15, false}, - SelectTestSpec{1, 128, 15, false}, - SelectTestSpec{20, 700, 1, true}, - SelectTestSpec{20, 700, 2, true}, - SelectTestSpec{20, 700, 3, true}, - SelectTestSpec{20, 700, 4, true}, - SelectTestSpec{20, 700, 5, true}, - SelectTestSpec{20, 700, 6, true}, - SelectTestSpec{20, 700, 7, true}, - SelectTestSpec{20, 700, 8, true}, - SelectTestSpec{20, 700, 9, true}, - SelectTestSpec{20, 700, 10, true, false}, - SelectTestSpec{20, 700, 11, true}, - SelectTestSpec{20, 700, 12, true}, - SelectTestSpec{20, 700, 16, true}, - SelectTestSpec{100, 1700, 17, true}, - SelectTestSpec{100, 1700, 31, true, false}, - SelectTestSpec{100, 1700, 32, false}, - SelectTestSpec{100, 1700, 33, false}, - SelectTestSpec{100, 1700, 63, false}, - SelectTestSpec{100, 1700, 64, false, false}, - SelectTestSpec{100, 1700, 65, false}, - SelectTestSpec{100, 1700, 255, true}, - SelectTestSpec{100, 1700, 256, true}, - SelectTestSpec{100, 1700, 511, false}, - SelectTestSpec{100, 1700, 512, true}, - SelectTestSpec{100, 1700, 1023, false, false}, - SelectTestSpec{100, 1700, 1024, true}, - SelectTestSpec{100, 1700, 1700, true}); - -auto inputs_random_largesize = testing::Values(SelectTestSpec{100, 100000, 1, true}, - SelectTestSpec{100, 100000, 2, true}, - SelectTestSpec{100, 100000, 3, true, false}, - SelectTestSpec{100, 100000, 7, true}, - SelectTestSpec{100, 100000, 16, true}, - SelectTestSpec{100, 100000, 31, true}, - SelectTestSpec{100, 100000, 32, true, false}, - SelectTestSpec{100, 100000, 60, true}, - SelectTestSpec{100, 100000, 100, true, false}, - SelectTestSpec{100, 100000, 200, true}, - SelectTestSpec{100000, 100, 100, false}, - SelectTestSpec{1, 100000000, 1, true}, - SelectTestSpec{1, 100000000, 16, false, false}, - SelectTestSpec{1, 100000000, 64, false}, - SelectTestSpec{1, 100000000, 128, true, false}, - SelectTestSpec{1, 100000000, 256, false, false}); - -auto inputs_random_largek = testing::Values(SelectTestSpec{100, 100000, 1000, true}, - SelectTestSpec{100, 100000, 2000, false}, - SelectTestSpec{100, 100000, 100000, true, false}, - SelectTestSpec{100, 100000, 2048, false}, - SelectTestSpec{100, 100000, 1237, true}); - -typedef SelectionTest::params_random> - ReferencedRandomFloatInt; -TEST_P(ReferencedRandomFloatInt, Run) { run(); } -INSTANTIATE_TEST_CASE_P(SelectionTest, - ReferencedRandomFloatInt, - testing::Combine(inputs_random_longlist, - testing::Values(knn::SelectKAlgo::RADIX_8_BITS, - knn::SelectKAlgo::RADIX_11_BITS, - knn::SelectKAlgo::WARP_SORT), - testing::Values(std::make_shared()))); - -typedef SelectionTest::params_random> - ReferencedRandomDoubleSizeT; -TEST_P(ReferencedRandomDoubleSizeT, Run) { run(); } -INSTANTIATE_TEST_CASE_P(SelectionTest, - ReferencedRandomDoubleSizeT, - testing::Combine(inputs_random_longlist, - testing::Values(knn::SelectKAlgo::RADIX_8_BITS, - knn::SelectKAlgo::RADIX_11_BITS, - knn::SelectKAlgo::WARP_SORT), - testing::Values(std::make_shared()))); - -typedef SelectionTest::params_random> - ReferencedRandomDoubleInt; -TEST_P(ReferencedRandomDoubleInt, LargeSize) { run(); } -INSTANTIATE_TEST_CASE_P(SelectionTest, - ReferencedRandomDoubleInt, - testing::Combine(inputs_random_largesize, - testing::Values(knn::SelectKAlgo::WARP_SORT), - testing::Values(std::make_shared()))); - -/** TODO: Fix test failure in RAFT CI - * - * SelectionTest/ReferencedRandomFloatSizeT.LargeK/0 - * Indicices do not match! ref[91628] = 131.359 != res[36504] = 158.438 - * Actual: false (actual=36504 != expected=91628 @38999; - * - * SelectionTest/ReferencedRandomFloatSizeT.LargeK/1 - * ERROR: ref[57977] = 58.9079 != res[21973] = 54.9354 - * Actual: false (actual=21973 != expected=57977 @107999; - * - */ -typedef SelectionTest::params_random> - ReferencedRandomFloatSizeT; -TEST_P(ReferencedRandomFloatSizeT, LargeK) { run(); } -INSTANTIATE_TEST_CASE_P(SelectionTest, - ReferencedRandomFloatSizeT, - testing::Combine(inputs_random_largek, - testing::Values(knn::SelectKAlgo::FAISS), - testing::Values(std::make_shared()))); -} // namespace raft::spatial::selection diff --git a/cpp/test/neighbors/tiled_knn.cu b/cpp/test/neighbors/tiled_knn.cu index a84c9749d7..f41a36dde3 100644 --- a/cpp/test/neighbors/tiled_knn.cu +++ b/cpp/test/neighbors/tiled_knn.cu @@ -27,7 +27,6 @@ #include #include #include // raft::neighbors::detail::brute_force_knn_impl -#include // raft::neighbors::detail::select_k #include @@ -128,15 +127,14 @@ class TiledKNNTest : public ::testing::TestWithParam { temp_dist = temp_row_major_dist.data(); } - raft::neighbors::detail::select_k(temp_dist, - nullptr, - num_queries, - num_db_vecs, - ref_distances_.data(), - ref_indices_.data(), - raft::distance::is_min_close(metric), - k_, - stream_); + matrix::select_k( + handle_, + raft::make_device_matrix_view(temp_dist, num_queries, num_db_vecs), + std::nullopt, + raft::make_device_matrix_view(ref_distances_.data(), params_.num_queries, params_.k), + raft::make_device_matrix_view(ref_indices_.data(), params_.num_queries, params_.k), + raft::distance::is_min_close(metric), + true); if ((params_.row_tiles == 0) && (params_.col_tiles == 0)) { std::vector input{database.data()};