From 6292fd7eba1f5c75806ee77058be6ee4ce598af2 Mon Sep 17 00:00:00 2001 From: Ben Frederickson Date: Wed, 29 Nov 2023 13:52:50 -0800 Subject: [PATCH] Remove selection_faiss 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. --- 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()};