From d534487e729f5ad54e4d6452284e6c555ea61b42 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Thu, 15 Feb 2024 01:33:40 -0800 Subject: [PATCH 1/7] initial commit --- .../detail/cagra/compute_distance.hpp | 25 ++++++++++++++----- 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp b/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp index 7f2e8b34cb..a94ee161d1 100644 --- a/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp +++ b/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp @@ -19,6 +19,7 @@ #include "device_common.hpp" #include "hashmap.hpp" +#include "raft/distance/distance_types.hpp" #include "utils.hpp" #include @@ -61,7 +62,8 @@ struct distance_op(); @@ -87,8 +89,13 @@ struct distance_op= dataset_dim) break; DISTANCE_T diff = query_buffer[device::swizzling(kv)]; - diff -= spatial::knn::detail::utils::mapping{}(dl_buff[e].data[v]); - norm2 += diff * diff; + if (metric == raft::distance::L2Expanded) { + diff -= spatial::knn::detail::utils::mapping{}(dl_buff[e].data[v]); + norm2 += diff * diff; + } else { + diff *= spatial::knn::detail::utils::mapping{}(dl_buff[e].data[v]); + norm2 += diff; + } } } } @@ -130,7 +137,8 @@ struct distance_op(); @@ -155,8 +163,13 @@ struct distance_op{}(dl_buff[e].data[v]); - norm2 += diff * diff; + if (metric == raft::distance::L2Expanded) { + diff -= spatial::knn::detail::utils::mapping{}(dl_buff[e].data[v]); + norm2 += diff * diff; + } else { + diff *= spatial::knn::detail::utils::mapping{}(dl_buff[e].data[v]); + norm2 += diff; + } } } } From c133c898afab3d0a70ca81b79af238180b7108c9 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Fri, 16 Feb 2024 04:47:31 -0800 Subject: [PATCH 2/7] change default use_raft in GpuClonerOptions and GpuDistance --- cpp/include/raft/neighbors/cagra_types.hpp | 3 ++ .../detail/cagra/compute_distance.hpp | 8 +++-- .../detail/cagra/search_multi_cta.cuh | 1 + .../cagra/search_multi_cta_kernel-ext.cuh | 2 ++ .../cagra/search_multi_cta_kernel-inl.cuh | 11 +++++-- .../detail/cagra/search_multi_kernel.cuh | 20 +++++++++---- .../detail/cagra/search_single_cta.cuh | 1 + .../cagra/search_single_cta_kernel-ext.cuh | 2 ++ .../cagra/search_single_cta_kernel-inl.cuh | 13 +++++--- .../cagra/search_multi_cta_00_generate.py | 1 + ...arch_multi_cta_float_uint32_dim1024_t32.cu | 1 + ...search_multi_cta_float_uint32_dim128_t8.cu | 1 + ...earch_multi_cta_float_uint32_dim256_t16.cu | 1 + ...earch_multi_cta_float_uint32_dim512_t32.cu | 1 + ...arch_multi_cta_float_uint64_dim1024_t32.cu | 1 + ...search_multi_cta_float_uint64_dim128_t8.cu | 1 + ...earch_multi_cta_float_uint64_dim256_t16.cu | 1 + ...earch_multi_cta_float_uint64_dim512_t32.cu | 1 + ...earch_multi_cta_half_uint32_dim1024_t32.cu | 1 + .../search_multi_cta_half_uint32_dim128_t8.cu | 1 + ...search_multi_cta_half_uint32_dim256_t16.cu | 1 + ...search_multi_cta_half_uint32_dim512_t32.cu | 1 + ...earch_multi_cta_half_uint64_dim1024_t32.cu | 1 + .../search_multi_cta_half_uint64_dim128_t8.cu | 1 + ...search_multi_cta_half_uint64_dim256_t16.cu | 1 + ...search_multi_cta_half_uint64_dim512_t32.cu | 1 + ...earch_multi_cta_int8_uint32_dim1024_t32.cu | 1 + .../search_multi_cta_int8_uint32_dim128_t8.cu | 1 + ...search_multi_cta_int8_uint32_dim256_t16.cu | 1 + ...search_multi_cta_int8_uint32_dim512_t32.cu | 1 + ...arch_multi_cta_uint8_uint32_dim1024_t32.cu | 1 + ...search_multi_cta_uint8_uint32_dim128_t8.cu | 1 + ...earch_multi_cta_uint8_uint32_dim256_t16.cu | 1 + ...earch_multi_cta_uint8_uint32_dim512_t32.cu | 1 + .../cagra/search_single_cta_00_generate.py | 1 + ...rch_single_cta_float_uint32_dim1024_t32.cu | 1 + ...earch_single_cta_float_uint32_dim128_t8.cu | 1 + ...arch_single_cta_float_uint32_dim256_t16.cu | 1 + ...arch_single_cta_float_uint32_dim512_t32.cu | 1 + ...rch_single_cta_float_uint64_dim1024_t32.cu | 1 + ...earch_single_cta_float_uint64_dim128_t8.cu | 1 + ...arch_single_cta_float_uint64_dim256_t16.cu | 1 + ...arch_single_cta_float_uint64_dim512_t32.cu | 1 + ...arch_single_cta_half_uint32_dim1024_t32.cu | 1 + ...search_single_cta_half_uint32_dim128_t8.cu | 1 + ...earch_single_cta_half_uint32_dim256_t16.cu | 1 + ...earch_single_cta_half_uint32_dim512_t32.cu | 1 + ...arch_single_cta_half_uint64_dim1024_t32.cu | 1 + ...search_single_cta_half_uint64_dim128_t8.cu | 1 + ...earch_single_cta_half_uint64_dim256_t16.cu | 1 + ...earch_single_cta_half_uint64_dim512_t32.cu | 1 + ...arch_single_cta_int8_uint32_dim1024_t32.cu | 1 + ...search_single_cta_int8_uint32_dim128_t8.cu | 1 + ...earch_single_cta_int8_uint32_dim256_t16.cu | 1 + ...earch_single_cta_int8_uint32_dim512_t32.cu | 1 + ...rch_single_cta_uint8_uint32_dim1024_t32.cu | 1 + ...earch_single_cta_uint8_uint32_dim128_t8.cu | 1 + ...arch_single_cta_uint8_uint32_dim256_t16.cu | 1 + ...arch_single_cta_uint8_uint32_dim512_t32.cu | 1 + .../ann_cagra/search_kernel_uint64_t.cuh | 2 ++ test.py | 30 +++++++++++++++++++ 61 files changed, 127 insertions(+), 16 deletions(-) create mode 100644 test.py diff --git a/cpp/include/raft/neighbors/cagra_types.hpp b/cpp/include/raft/neighbors/cagra_types.hpp index 00c363b377..a59b101c9d 100644 --- a/cpp/include/raft/neighbors/cagra_types.hpp +++ b/cpp/include/raft/neighbors/cagra_types.hpp @@ -116,6 +116,9 @@ struct search_params : ann::search_params { uint32_t num_random_samplings = 1; /** Bit mask used for initial random seed node selection. */ uint64_t rand_xor_mask = 0x128394; + + /** Distance metric */ + distance::DistanceType metric = distance::L2Expanded; }; static_assert(std::is_aggregate_v); diff --git a/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp b/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp index a94ee161d1..72181283e3 100644 --- a/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp +++ b/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp @@ -201,6 +201,7 @@ _RAFT_DEVICE void compute_distance_to_random_nodes( const uint32_t num_seeds, INDEX_T* const visited_hash_ptr, const uint32_t hash_bitlen, + raft::distance::DistanceType metric, const uint32_t block_id = 0, const uint32_t num_blocks = 1) { @@ -228,7 +229,7 @@ _RAFT_DEVICE void compute_distance_to_random_nodes( } } - const auto norm2 = dist_op(dataset_ptr + dataset_ld * seed_index, dataset_dim, valid_i); + const auto norm2 = dist_op(dataset_ptr + dataset_ld * seed_index, dataset_dim, valid_i, metric); if (valid_i && (norm2 < best_norm2_team_local)) { best_norm2_team_local = norm2; @@ -272,7 +273,8 @@ _RAFT_DEVICE void compute_distance_to_child_nodes(INDEX_T* const result_child_in const std::uint32_t hash_bitlen, const INDEX_T* const parent_indices, const INDEX_T* const internal_topk_list, - const std::uint32_t search_width) + const std::uint32_t search_width, + raft::distance::DistanceType metric) { constexpr INDEX_T index_msb_1_mask = utils::gen_index_msb_1_mask::value; const INDEX_T invalid_index = utils::get_max_value(); @@ -315,7 +317,7 @@ _RAFT_DEVICE void compute_distance_to_child_nodes(INDEX_T* const result_child_in if (valid_i) { child_id = result_child_indices_ptr[i]; } DISTANCE_T norm2 = - dist_op(dataset_ptr + child_id * dataset_ld, dataset_dim, child_id != invalid_index); + dist_op(dataset_ptr + child_id * dataset_ld, dataset_dim, child_id != invalid_index, metric); // Store the distance const unsigned lane_id = threadIdx.x % TEAM_SIZE; diff --git a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh index 010b0a6f80..a0dd7f8dd3 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh @@ -230,6 +230,7 @@ struct search : public search_plan_implmetric, stream); RAFT_CUDA_TRY(cudaPeekAtLastError()); diff --git a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-ext.cuh b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-ext.cuh index 7a5ad17460..a1ad6320f9 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-ext.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-ext.cuh @@ -54,6 +54,7 @@ void select_and_run(raft::device_matrix_view= num_pickup INDEX_T* const visited_hashmap_ptr, // [num_queries, 1 << bitlen] - const std::uint32_t hash_bitlen) + const std::uint32_t hash_bitlen, + distance::DistanceType metric) { const auto ldb = hashmap::get_size(hash_bitlen); const auto global_team_index = (blockIdx.x * blockDim.x + threadIdx.x) / TEAM_SIZE; @@ -138,7 +139,7 @@ RAFT_KERNEL random_pickup_kernel(const DATA_T* const dataset_ptr, // [dataset_s seed_index = device::xorshift64((global_team_index ^ rand_xor_mask) * (i + 1)) % dataset_size; } - const auto norm2 = dist_op(dataset_ptr + (dataset_ld * seed_index), dataset_dim, true); + const auto norm2 = dist_op(dataset_ptr + (dataset_ld * seed_index), dataset_dim, true, metric); if (norm2 < best_norm2_team_local) { best_norm2_team_local = norm2; @@ -181,6 +182,7 @@ void random_pickup(const DATA_T* const dataset_ptr, // [dataset_size, dataset_d const std::size_t ldr, // (*) ldr >= num_pickup INDEX_T* const visited_hashmap_ptr, // [num_queries, 1 << bitlen] const std::uint32_t hash_bitlen, + distance::DistanceType metric, cudaStream_t const cuda_stream = 0) { const auto block_size = 256u; @@ -207,7 +209,8 @@ void random_pickup(const DATA_T* const dataset_ptr, // [dataset_size, dataset_d result_distances_ptr, ldr, visited_hashmap_ptr, - hash_bitlen); + hash_bitlen, + metric); } template @@ -334,7 +337,8 @@ RAFT_KERNEL compute_distance_to_child_nodes_kernel( INDEX_T* const result_indices_ptr, // [num_queries, ldd] DISTANCE_T* const result_distances_ptr, // [num_queries, ldd] const std::uint32_t ldd, // (*) ldd >= search_width * graph_degree - SAMPLE_FILTER_T sample_filter) + SAMPLE_FILTER_T sample_filter, + distance::DistanceType metric) { const uint32_t ldb = hashmap::get_size(hash_bitlen); const auto tid = threadIdx.x + blockDim.x * blockIdx.x; @@ -381,7 +385,7 @@ RAFT_KERNEL compute_distance_to_child_nodes_kernel( visited_hashmap_ptr + (ldb * blockIdx.y), hash_bitlen, child_id); const auto norm2 = - dist_op(dataset_ptr + (dataset_ld * child_id), dataset_dim, compute_distance_flag); + dist_op(dataset_ptr + (dataset_ld * child_id), dataset_dim, compute_distance_flag, metric); if (compute_distance_flag) { if (threadIdx.x % TEAM_SIZE == 0) { @@ -430,6 +434,7 @@ void compute_distance_to_child_nodes( DISTANCE_T* const result_distances_ptr, // [num_queries, ldd] const std::uint32_t ldd, // (*) ldd >= search_width * graph_degree SAMPLE_FILTER_T sample_filter, + distance::DistanceType metric, cudaStream_t cuda_stream = 0) { const auto block_size = 128; @@ -460,7 +465,8 @@ void compute_distance_to_child_nodes( result_indices_ptr, result_distances_ptr, ldd, - sample_filter); + sample_filter, + metric); } template @@ -844,6 +850,7 @@ struct search : search_plan_impl { result_buffer_allocation_size, hashmap.data(), hash_bitlen, + this->metric, stream); unsigned iter = 0; @@ -916,6 +923,7 @@ struct search : search_plan_impl { result_distances.data() + itopk_size, result_buffer_allocation_size, sample_filter, + this->metric, stream); iter++; diff --git a/cpp/include/raft/neighbors/detail/cagra/search_single_cta.cuh b/cpp/include/raft/neighbors/detail/cagra/search_single_cta.cuh index 0b4fc2d47b..ba001901cf 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_single_cta.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_single_cta.cuh @@ -241,6 +241,7 @@ struct search : search_plan_impl { min_iterations, max_iterations, sample_filter, + this->metric, stream); } }; diff --git a/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-ext.cuh b/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-ext.cuh index fef060ffee..7f023668fa 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-ext.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-ext.cuh @@ -56,6 +56,7 @@ void select_and_run( // raft::resources const& res, size_t min_iterations, size_t max_iterations, SAMPLE_FILTER_T sample_filter, + distance::DistanceType metric, cudaStream_t stream) RAFT_EXPLICIT; #endif // RAFT_EXPLICIT_INSTANTIATE_ONLY @@ -88,6 +89,7 @@ void select_and_run( // raft::resources const& res, size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 80b5b343b2..10321e215d 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -487,7 +487,8 @@ __launch_bounds__(1024, 1) RAFT_KERNEL const std::uint32_t hash_bitlen, const std::uint32_t small_hash_bitlen, const std::uint32_t small_hash_reset_interval, - SAMPLE_FILTER_T sample_filter) + SAMPLE_FILTER_T sample_filter, + raft::distance::DistanceType metric) { using LOAD_T = device::LOAD_128BIT_T; const auto query_id = blockIdx.y; @@ -580,7 +581,8 @@ __launch_bounds__(1024, 1) RAFT_KERNEL local_seed_ptr, num_seeds, local_visited_hashmap_ptr, - hash_bitlen); + hash_bitlen, + metric); __syncthreads(); _CLK_REC(clk_compute_1st_distance); @@ -719,7 +721,8 @@ __launch_bounds__(1024, 1) RAFT_KERNEL hash_bitlen, parent_list_buffer, result_indices_buffer, - search_width); + search_width, + metric); __syncthreads(); _CLK_REC(clk_compute_distance); @@ -918,6 +921,7 @@ void select_and_run( // raft::resources const& res, size_t min_iterations, size_t max_iterations, SAMPLE_FILTER_T sample_filter, + distance::DistanceType metric, cudaStream_t stream) { auto kernel = @@ -958,7 +962,8 @@ void select_and_run( // raft::resources const& res, hash_bitlen, small_hash_bitlen, small_hash_reset_interval, - sample_filter); + sample_filter, + metric); RAFT_CUDA_TRY(cudaPeekAtLastError()); } } // namespace single_cta_search diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_00_generate.py b/cpp/src/neighbors/detail/cagra/search_multi_cta_00_generate.py index 6f8766c86b..8c1e13283f 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_00_generate.py +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_00_generate.py @@ -70,6 +70,7 @@ size_t min_iterations, \\ size_t max_iterations, \\ SAMPLE_FILTER_T sample_filter, \\ + distance::DistanceType metric, \\ cudaStream_t stream); """ diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim1024_t32.cu index 1a3b2284bd..4f7b3a5935 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim1024_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim128_t8.cu index 36e86d9ed6..5e88018fa2 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim128_t8.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim256_t16.cu index 6f1af2d93f..22659f2098 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim256_t16.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim512_t32.cu index 1279f8e415..89c6b630c2 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim512_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim1024_t32.cu index 0dabff0df5..34c2b016c9 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim1024_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim128_t8.cu index 72bb74cdb8..7ab3cc6446 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim128_t8.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim256_t16.cu index dceea10b5d..7f93961b9d 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim256_t16.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim512_t32.cu index acb8bd6a12..bf12198b4c 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_float_uint64_dim512_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim1024_t32.cu index fa89bca45f..723bbc56fd 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim1024_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim128_t8.cu index 645ca61ff5..1e1e1bfaae 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim128_t8.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim256_t16.cu index 41b6f9b420..d16d713c06 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim256_t16.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim512_t32.cu index 38f0ac3b04..23598ec29d 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint32_dim512_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim1024_t32.cu index c462a9d359..9691753060 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim1024_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim128_t8.cu index f5b2874e20..21cdd8c04e 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim128_t8.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim256_t16.cu index 0b01428b86..7dcc72b8f1 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim256_t16.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim512_t32.cu index 70228a129d..ff67019b5a 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_half_uint64_dim512_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim1024_t32.cu index 0254f09ff0..3568cec680 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim1024_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim128_t8.cu index 2b67e7e968..71261268e8 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim128_t8.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim256_t16.cu index 17d6722e58..3f5beefb16 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim256_t16.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim512_t32.cu index 38f02812e2..844bfa7ff2 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim512_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim1024_t32.cu index fa111196c6..01568627ad 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim1024_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim128_t8.cu index 1ef3c28aa3..46e2a21656 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim128_t8.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim256_t16.cu index d26cb44843..e784fc8e53 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim256_t16.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim512_t32.cu index 4d4322f261..65306f7549 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim512_t32.cu @@ -56,6 +56,7 @@ namespace raft::neighbors::cagra::detail::multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_00_generate.py b/cpp/src/neighbors/detail/cagra/search_single_cta_00_generate.py index 1515f43134..2e3115c432 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_00_generate.py +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_00_generate.py @@ -71,6 +71,7 @@ size_t min_iterations, \\ size_t max_iterations, \\ SAMPLE_FILTER_T sample_filter, \\ + distance::DistanceType metric, \\ cudaStream_t stream); """ diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim1024_t32.cu index b8c23103ba..e201e0c179 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim1024_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim128_t8.cu index 8ab1897119..af7d4a0695 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim128_t8.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim256_t16.cu index 9fd36b4cb9..3e1743dda5 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim256_t16.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim512_t32.cu index a9ee2c864b..11a622a928 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint32_dim512_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim1024_t32.cu index dadc574b65..81759d998b 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim1024_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim128_t8.cu index 30e043f47e..56cb1996e3 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim128_t8.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim256_t16.cu index 089e4c930f..4829f61922 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim256_t16.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim512_t32.cu index 3e8ffb8bf8..23b7dc61f5 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_float_uint64_dim512_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim1024_t32.cu index 29e7bfa250..f0e6fc635f 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim1024_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim128_t8.cu index a004f900d0..64326d590a 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim128_t8.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim256_t16.cu index 549849b21d..86ea5fd508 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim256_t16.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim512_t32.cu index 3825f572f7..64c09b4b1b 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint32_dim512_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim1024_t32.cu index 31d83f443b..f5346862f3 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim1024_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim128_t8.cu index 3493ab294c..fb8cbf577f 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim128_t8.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim256_t16.cu index 6e09709994..b2993a4704 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim256_t16.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim512_t32.cu index 4bc0158f7e..def38359b1 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_half_uint64_dim512_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim1024_t32.cu index 279587738e..9e9445107e 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim1024_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim128_t8.cu index ef127d3f7d..dc7747d353 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim128_t8.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim256_t16.cu index 7fcfdcc28e..c5d7c7e412 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim256_t16.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim512_t32.cu index a6c606d99b..2b0c1edff2 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim512_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim1024_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim1024_t32.cu index 0b8be56614..c6e1ac120f 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim1024_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim1024_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim128_t8.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim128_t8.cu index 4c193b9408..a3d5810776 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim128_t8.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim128_t8.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim256_t16.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim256_t16.cu index bdf16d2f03..5d0446de64 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim256_t16.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim256_t16.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim512_t32.cu b/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim512_t32.cu index 93624df4aa..6a3296bbfb 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim512_t32.cu +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim512_t32.cu @@ -57,6 +57,7 @@ namespace raft::neighbors::cagra::detail::single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/cpp/test/neighbors/ann_cagra/search_kernel_uint64_t.cuh b/cpp/test/neighbors/ann_cagra/search_kernel_uint64_t.cuh index 175e4ef483..d7cf64a15d 100644 --- a/cpp/test/neighbors/ann_cagra/search_kernel_uint64_t.cuh +++ b/cpp/test/neighbors/ann_cagra/search_kernel_uint64_t.cuh @@ -48,6 +48,7 @@ namespace multi_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_kernel_selection( @@ -92,6 +93,7 @@ namespace single_cta_search { size_t min_iterations, \ size_t max_iterations, \ SAMPLE_FILTER_T sample_filter, \ + distance::DistanceType metric, \ cudaStream_t stream); instantiate_single_cta_select_and_run( diff --git a/test.py b/test.py new file mode 100644 index 0000000000..60b7b92a01 --- /dev/null +++ b/test.py @@ -0,0 +1,30 @@ +import numpy as np +import faiss +rs = np.random.RandomState(123) + +xq = rs.rand(50000, 256).astype('float32') +xb = rs.rand(500000, 256).astype('float32') +res = faiss.StandardGpuResources() +quantizer = faiss.IndexFlatL2(256) # This is an index that performs a brute-force L2 search +index = faiss.IndexIVFFlat(quantizer, 256, 10000, faiss.METRIC_L2) +index.train(xb) +index_gpu_classical = faiss.index_cpu_to_gpu(res, 0, index) +index_gpu_classical.add(xb) +Dref, Iref = index_gpu_classical.search(xq, 10000) +# index = faiss.index_factory(256, "Flat") +index = faiss.IndexIVFFlat(quantizer, 256, 10000, faiss.METRIC_L2) +co = faiss.GpuMultipleClonerOptions() +print(co.use_raft) +co.useFloat16 = True +index_gpu_fp16 = faiss.index_cpu_to_gpu(res, 0, index, co) +index_gpu_fp16.add(xb) +D, I = index_gpu_fp16.search(xq, 100) +(I != Iref).sum() / I.size +# index = faiss.index_factory(256, "Flat") +co = faiss.GpuMultipleClonerOptions() +co.use_raft = True +index_gpu_raft = faiss.index_cpu_to_gpu(res, 0, index, co) +index_gpu_raft.add(xb) +D, I = index_gpu_raft.search(xq, 10) +D, I = index_gpu_raft.search(xq, 200) +print(D, I) From ec2f1f598cd943d620e2c77ef49781f17ddec08f Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Tue, 20 Feb 2024 06:19:26 -0800 Subject: [PATCH 3/7] nit --- cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp b/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp index 72181283e3..a8dd6eddc3 100644 --- a/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp +++ b/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp @@ -15,11 +15,11 @@ */ #pragma once +#include #include #include "device_common.hpp" #include "hashmap.hpp" -#include "raft/distance/distance_types.hpp" #include "utils.hpp" #include From a374a51443041e278bced4fce6c896397bb2a740 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 21 Feb 2024 10:34:40 -0800 Subject: [PATCH 4/7] more changes --- cpp/include/raft/neighbors/cagra_types.hpp | 3 --- .../raft/neighbors/detail/cagra/cagra_search.cuh | 2 +- .../neighbors/detail/cagra/compute_distance.hpp | 4 ++-- .../raft/neighbors/detail/cagra/factory.cuh | 5 +++-- .../detail/cagra/search_multi_cta_kernel-inl.cuh | 2 ++ .../raft/neighbors/detail/cagra/search_plan.cuh | 11 +++++++---- cpp/test/neighbors/ann_cagra.cuh | 14 +++++++------- 7 files changed, 22 insertions(+), 19 deletions(-) diff --git a/cpp/include/raft/neighbors/cagra_types.hpp b/cpp/include/raft/neighbors/cagra_types.hpp index d59e673b19..cc76d70fbf 100644 --- a/cpp/include/raft/neighbors/cagra_types.hpp +++ b/cpp/include/raft/neighbors/cagra_types.hpp @@ -116,9 +116,6 @@ struct search_params : ann::search_params { uint32_t num_random_samplings = 1; /** Bit mask used for initial random seed node selection. */ uint64_t rand_xor_mask = 0x128394; - - /** Distance metric */ - distance::DistanceType metric = distance::L2Expanded; }; static_assert(std::is_aggregate_v); diff --git a/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh b/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh index 40cc7c76fb..7933ed67b9 100644 --- a/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh @@ -129,7 +129,7 @@ void search_main(raft::resources const& res, using CagraSampleFilterT_s = typename CagraSampleFilterT_Selector::type; std::unique_ptr> plan = factory::create( - res, params, index.dim(), index.graph_degree(), topk); + res, params, index.dim(), index.graph_degree(), topk, index.metric()); plan->check(topk); diff --git a/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp b/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp index a8dd6eddc3..2fb62097b7 100644 --- a/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp +++ b/cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp @@ -94,7 +94,7 @@ struct distance_op{}(dl_buff[e].data[v]); - norm2 += diff; + norm2 -= diff; } } } @@ -168,7 +168,7 @@ struct distance_op{}(dl_buff[e].data[v]); - norm2 += diff; + norm2 -= diff; } } } diff --git a/cpp/include/raft/neighbors/detail/cagra/factory.cuh b/cpp/include/raft/neighbors/detail/cagra/factory.cuh index 0002dd8b2a..4b82fde312 100644 --- a/cpp/include/raft/neighbors/detail/cagra/factory.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/factory.cuh @@ -38,9 +38,10 @@ class factory { search_params const& params, int64_t dim, int64_t graph_degree, - uint32_t topk) + uint32_t topk, + distance::DistanceType metric) { - search_plan_impl_base plan(params, dim, graph_degree, topk); + search_plan_impl_base plan(params, dim, graph_degree, topk, metric); switch (plan.dataset_block_dim) { case 128: switch (plan.team_size) { diff --git a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index 07eb8afc89..2eb2ac1228 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -34,6 +34,7 @@ #include "compute_distance.hpp" #include "device_common.hpp" #include "hashmap.hpp" +#include "raft/distance/distance_types.hpp" #include "search_plan.cuh" #include "topk_for_cagra/topk_core.cuh" // TODO replace with raft topk if possible #include "utils.hpp" @@ -329,6 +330,7 @@ __launch_bounds__(1024, 1) RAFT_KERNEL search_kernel( result_distances_buffer[i] = utils::get_max_value(); result_indices_buffer[i] = invalid_index; } + if (metric == distance::InnerProduct) result_distances_buffer[i] *= -1; } __syncthreads(); diff --git a/cpp/include/raft/neighbors/detail/cagra/search_plan.cuh b/cpp/include/raft/neighbors/detail/cagra/search_plan.cuh index 271a1f4955..7e3202a00f 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_plan.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_plan.cuh @@ -23,6 +23,7 @@ #include #include +#include #include #include @@ -33,8 +34,9 @@ struct search_plan_impl_base : public search_params { int64_t dim; int64_t graph_degree; uint32_t topk; - search_plan_impl_base(search_params params, int64_t dim, int64_t graph_degree, uint32_t topk) - : search_params(params), dim(dim), graph_degree(graph_degree), topk(topk) + distance::DistanceType metric; + search_plan_impl_base(search_params params, int64_t dim, int64_t graph_degree, uint32_t topk, distance::DistanceType metric = distance::L2Expanded) + : search_params(params), dim(dim), graph_degree(graph_degree), topk(topk), metric(metric) { set_dataset_block_and_team_size(dim); if (algo == search_algo::AUTO) { @@ -91,8 +93,9 @@ struct search_plan_impl : public search_plan_impl_base { search_params params, int64_t dim, int64_t graph_degree, - uint32_t topk) - : search_plan_impl_base(params, dim, graph_degree, topk), + uint32_t topk, + distance::DistanceType metric = distance::L2Expanded) + : search_plan_impl_base(params, dim, graph_degree, topk, metric), hashmap(0, resource::get_cuda_stream(res)), num_executed_iterations(0, resource::get_cuda_stream(res)), dev_seed(0, resource::get_cuda_stream(res)), diff --git a/cpp/test/neighbors/ann_cagra.cuh b/cpp/test/neighbors/ann_cagra.cuh index 296a5f07fc..492856adec 100644 --- a/cpp/test/neighbors/ann_cagra.cuh +++ b/cpp/test/neighbors/ann_cagra.cuh @@ -738,7 +738,7 @@ inline std::vector generate_inputs() {0}, {256}, {1}, - {raft::distance::DistanceType::L2Expanded}, + {raft::distance::DistanceType::L2Expanded, raft::distance::DistanceType::InnerProduct}, {false}, {true}, {0.995}); @@ -754,7 +754,7 @@ inline std::vector generate_inputs() {0}, {256}, {1}, - {raft::distance::DistanceType::L2Expanded}, + {raft::distance::DistanceType::L2Expanded, raft::distance::DistanceType::InnerProduct}, {false}, {true}, {99. / 100} @@ -773,7 +773,7 @@ inline std::vector generate_inputs() {0}, {64}, {1}, - {raft::distance::DistanceType::L2Expanded}, + {raft::distance::DistanceType::L2Expanded, raft::distance::DistanceType::InnerProduct}, {false}, {true}, {0.995}); @@ -789,7 +789,7 @@ inline std::vector generate_inputs() {0, 4, 8, 16, 32}, // team_size {64}, {1}, - {raft::distance::DistanceType::L2Expanded}, + {raft::distance::DistanceType::L2Expanded, raft::distance::DistanceType::InnerProduct}, {false}, {false}, {0.995}); @@ -806,7 +806,7 @@ inline std::vector generate_inputs() {0}, // team_size {32, 64, 128, 256, 512, 768}, {1}, - {raft::distance::DistanceType::L2Expanded}, + {raft::distance::DistanceType::L2Expanded, raft::distance::DistanceType::InnerProduct}, {false}, {true}, {0.995}); @@ -823,7 +823,7 @@ inline std::vector generate_inputs() {0}, // team_size {64}, {1}, - {raft::distance::DistanceType::L2Expanded}, + {raft::distance::DistanceType::L2Expanded, raft::distance::DistanceType::InnerProduct}, {false, true}, {false}, {0.995}); @@ -840,7 +840,7 @@ inline std::vector generate_inputs() {0}, {4096}, // itopk_size {1}, - {raft::distance::DistanceType::L2Expanded}, + {raft::distance::DistanceType::L2Expanded, raft::distance::DistanceType::InnerProduct}, {false}, {false}, {0.995}); From 7b739c35b72642d2dd57c73c0566f63f8cb0ac01 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Mon, 26 Feb 2024 10:22:01 -0800 Subject: [PATCH 5/7] update search kernel --- .../detail/cagra/search_multi_cta_kernel-inl.cuh | 9 +++++++-- .../detail/cagra/search_single_cta_kernel-inl.cuh | 7 ++++++- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index 2eb2ac1228..6dc0f3be92 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -330,7 +330,6 @@ __launch_bounds__(1024, 1) RAFT_KERNEL search_kernel( result_distances_buffer[i] = utils::get_max_value(); result_indices_buffer[i] = invalid_index; } - if (metric == distance::InnerProduct) result_distances_buffer[i] *= -1; } __syncthreads(); @@ -343,7 +342,13 @@ __launch_bounds__(1024, 1) RAFT_KERNEL search_kernel( for (uint32_t i = threadIdx.x; i < itopk_size; i += blockDim.x) { uint32_t j = i + (itopk_size * (cta_id + (num_cta_per_query * query_id))); - if (result_distances_ptr != nullptr) { result_distances_ptr[j] = result_distances_buffer[i]; } + const INDEX_T invalid_index = utils::get_max_value(); + + if (result_distances_ptr != nullptr) { + if (metric == distance::InnerProduct && result_indices_buffer[i] != invalid_index) { + result_distances_ptr[j] = -result_distances_buffer[i]; + } else { + result_distances_ptr[j] = result_distances_buffer[i]; }} constexpr INDEX_T index_msb_1_mask = utils::gen_index_msb_1_mask::value; result_indices_ptr[j] = diff --git a/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 10321e215d..1d555c7259 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -781,7 +781,12 @@ __launch_bounds__(1024, 1) RAFT_KERNEL unsigned j = i + (top_k * query_id); unsigned ii = i; if (TOPK_BY_BITONIC_SORT) { ii = device::swizzling(i); } - if (result_distances_ptr != nullptr) { result_distances_ptr[j] = result_distances_buffer[ii]; } + const INDEX_T invalid_index = utils::get_max_value(); + if (result_distances_ptr != nullptr) { + if (metric == distance::InnerProduct && result_indices_buffer[ii] != invalid_index) { + result_distances_ptr[j] = -result_distances_buffer[ii]; + } else { + result_distances_ptr[j] = result_distances_buffer[ii]; }} constexpr INDEX_T index_msb_1_mask = utils::gen_index_msb_1_mask::value; result_indices_ptr[j] = From 28eb33aa5b82fa22548988cdabfaea921932ac32 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Thu, 29 Feb 2024 13:33:07 -0800 Subject: [PATCH 6/7] dbg: --- .../raft/neighbors/detail/cagra/cagra_search.cuh | 5 +++++ cpp/include/raft/neighbors/detail/cagra/factory.cuh | 6 +++--- .../neighbors/detail/cagra/search_multi_cta.cuh | 6 ++++-- .../neighbors/detail/cagra/search_multi_kernel.cuh | 6 ++++-- .../neighbors/detail/cagra/search_single_cta.cuh | 6 ++++-- cpp/test/neighbors/ann_cagra.cuh | 13 +++++++++++++ 6 files changed, 33 insertions(+), 9 deletions(-) diff --git a/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh b/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh index 7933ed67b9..98ed741188 100644 --- a/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh @@ -29,6 +29,8 @@ #include #include "factory.cuh" +#include "raft/distance/distance_types.hpp" +#include "raft/util/cudart_utils.hpp" #include "search_plan.cuh" #include "search_single_cta.cuh" @@ -171,10 +173,12 @@ void search_main(raft::resources const& res, _num_executed_iterations, topk, set_offset(sample_filter, qid)); + raft::print_device_vector("topk_distances_ptr", _topk_distances_ptr, 10, std::cout); } static_assert(std::is_same_v, "only float distances are supported at the moment"); + if (index.metric() != distance::InnerProduct) { float* dist_out = distances.data_handle(); const DistanceT* dist_in = distances.data_handle(); // We're converting the data from T to DistanceT during distance computation @@ -188,6 +192,7 @@ void search_main(raft::resources const& res, distances.extent(1), kScale, resource::get_cuda_stream(res)); + } } /** @} */ // end group cagra diff --git a/cpp/include/raft/neighbors/detail/cagra/factory.cuh b/cpp/include/raft/neighbors/detail/cagra/factory.cuh index 4b82fde312..d5a9aaa242 100644 --- a/cpp/include/raft/neighbors/detail/cagra/factory.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/factory.cuh @@ -75,17 +75,17 @@ class factory { return std::unique_ptr>( new single_cta_search:: search( - res, plan, plan.dim, plan.graph_degree, plan.topk)); + res, plan, plan.dim, plan.graph_degree, plan.topk, plan.metric)); } else if (plan.algo == search_algo::MULTI_CTA) { return std::unique_ptr>( new multi_cta_search:: search( - res, plan, plan.dim, plan.graph_degree, plan.topk)); + res, plan, plan.dim, plan.graph_degree, plan.topk, plan.metric)); } else { return std::unique_ptr>( new multi_kernel_search:: search( - res, plan, plan.dim, plan.graph_degree, plan.topk)); + res, plan, plan.dim, plan.graph_degree, plan.topk, plan.metric)); } } }; diff --git a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh index a0dd7f8dd3..7dfd874d7c 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta.cuh @@ -33,6 +33,7 @@ #include "compute_distance.hpp" #include "device_common.hpp" #include "hashmap.hpp" +#include "raft/distance/distance_types.hpp" #include "search_multi_cta_kernel.cuh" #include "search_plan.cuh" #include "topk_for_cagra/topk_core.cuh" // TODO replace with raft topk if possible @@ -95,9 +96,10 @@ struct search : public search_plan_impl( - res, params, dim, graph_degree, topk), + res, params, dim, graph_degree, topk, metric), intermediate_indices(0, resource::get_cuda_stream(res)), intermediate_distances(0, resource::get_cuda_stream(res)), topk_workspace(0, resource::get_cuda_stream(res)) diff --git a/cpp/include/raft/neighbors/detail/cagra/search_multi_kernel.cuh b/cpp/include/raft/neighbors/detail/cagra/search_multi_kernel.cuh index 4e947852db..2dbe8abbf0 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_multi_kernel.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_multi_kernel.cuh @@ -33,6 +33,7 @@ #include "compute_distance.hpp" #include "device_common.hpp" #include "hashmap.hpp" +#include "raft/distance/distance_types.hpp" #include "search_plan.cuh" #include "topk_for_cagra/topk_core.cuh" //todo replace with raft kernel #include "utils.hpp" @@ -670,9 +671,10 @@ struct search : search_plan_impl { search_params params, int64_t dim, int64_t graph_degree, - uint32_t topk) + uint32_t topk, + distance::DistanceType metric) : search_plan_impl( - res, params, dim, graph_degree, topk), + res, params, dim, graph_degree, topk, metric), result_indices(0, resource::get_cuda_stream(res)), result_distances(0, resource::get_cuda_stream(res)), parent_node_list(0, resource::get_cuda_stream(res)), diff --git a/cpp/include/raft/neighbors/detail/cagra/search_single_cta.cuh b/cpp/include/raft/neighbors/detail/cagra/search_single_cta.cuh index ba001901cf..1570374ee5 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_single_cta.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_single_cta.cuh @@ -33,6 +33,7 @@ #include "compute_distance.hpp" #include "device_common.hpp" #include "hashmap.hpp" +#include "raft/distance/distance_types.hpp" #include "search_plan.cuh" #include "search_single_cta_kernel.cuh" #include "topk_by_radix.cuh" @@ -91,9 +92,10 @@ struct search : search_plan_impl { search_params params, int64_t dim, int64_t graph_degree, - uint32_t topk) + uint32_t topk, + distance::DistanceType metric) : search_plan_impl( - res, params, dim, graph_degree, topk) + res, params, dim, graph_degree, topk, metric) { set_params(res); } diff --git a/cpp/test/neighbors/ann_cagra.cuh b/cpp/test/neighbors/ann_cagra.cuh index 492856adec..2f1abdc2b0 100644 --- a/cpp/test/neighbors/ann_cagra.cuh +++ b/cpp/test/neighbors/ann_cagra.cuh @@ -15,6 +15,7 @@ */ #pragma once +#include "raft/neighbors/cagra_types.hpp" #undef RAFT_EXPLICIT_INSTANTIATE_ONLY // Search with filter instantiation #include "../test_utils.cuh" @@ -222,6 +223,9 @@ class AnnCagraTest : public ::testing::TestWithParam { protected: void testCagra() { + // TODO (tarang-jain): remove when NN Descent index building support InnerProduct + if (ps.metric == distance::InnerProduct && ps.build_algo == graph_build_algo::NN_DESCENT) GTEST_SKIP(); + size_t queries_size = ps.n_queries * ps.k; std::vector indices_Cagra(queries_size); std::vector indices_naive(queries_size); @@ -254,6 +258,7 @@ class AnnCagraTest : public ::testing::TestWithParam { cagra::index_params index_params; index_params.metric = ps.metric; // Note: currently ony the cagra::index_params metric is // not used for knn_graph building. + RAFT_LOG_INFO("index_params.metric %d", index_params.metric); index_params.build_algo = ps.build_algo; cagra::search_params search_params; search_params.algo = ps.algo; @@ -369,6 +374,8 @@ class AnnCagraSortTest : public ::testing::TestWithParam { protected: void testCagraSort() { + // TODO (tarang-jain): remove when NN Descent index building support InnerProduct + if (ps.metric == distance::InnerProduct && ps.build_algo == graph_build_algo::NN_DESCENT) GTEST_SKIP(); { // Step 1: Build a sorted KNN graph by CAGRA knn build auto database_view = raft::make_device_matrix_view( @@ -454,6 +461,9 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { protected: void testCagraFilter() { + // TODO (tarang-jain): remove when NN Descent index building support InnerProduct + if (ps.metric == distance::InnerProduct && ps.build_algo == graph_build_algo::NN_DESCENT) GTEST_SKIP(); + size_t queries_size = ps.n_queries * ps.k; std::vector indices_Cagra(queries_size); std::vector indices_naive(queries_size); @@ -574,6 +584,9 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { void testCagraRemoved() { + // TODO (tarang-jain): remove when NN Descent index building support InnerProduct + if (ps.metric == distance::InnerProduct && ps.build_algo == graph_build_algo::NN_DESCENT) GTEST_SKIP(); + size_t queries_size = ps.n_queries * ps.k; std::vector indices_Cagra(queries_size); std::vector indices_naive(queries_size); From a11bfdaf1c0fa00040802524aade50354af5cc39 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Mon, 4 Mar 2024 11:59:53 -0800 Subject: [PATCH 7/7] rm test.py --- test.py | 30 ------------------------------ 1 file changed, 30 deletions(-) delete mode 100644 test.py diff --git a/test.py b/test.py deleted file mode 100644 index 60b7b92a01..0000000000 --- a/test.py +++ /dev/null @@ -1,30 +0,0 @@ -import numpy as np -import faiss -rs = np.random.RandomState(123) - -xq = rs.rand(50000, 256).astype('float32') -xb = rs.rand(500000, 256).astype('float32') -res = faiss.StandardGpuResources() -quantizer = faiss.IndexFlatL2(256) # This is an index that performs a brute-force L2 search -index = faiss.IndexIVFFlat(quantizer, 256, 10000, faiss.METRIC_L2) -index.train(xb) -index_gpu_classical = faiss.index_cpu_to_gpu(res, 0, index) -index_gpu_classical.add(xb) -Dref, Iref = index_gpu_classical.search(xq, 10000) -# index = faiss.index_factory(256, "Flat") -index = faiss.IndexIVFFlat(quantizer, 256, 10000, faiss.METRIC_L2) -co = faiss.GpuMultipleClonerOptions() -print(co.use_raft) -co.useFloat16 = True -index_gpu_fp16 = faiss.index_cpu_to_gpu(res, 0, index, co) -index_gpu_fp16.add(xb) -D, I = index_gpu_fp16.search(xq, 100) -(I != Iref).sum() / I.size -# index = faiss.index_factory(256, "Flat") -co = faiss.GpuMultipleClonerOptions() -co.use_raft = True -index_gpu_raft = faiss.index_cpu_to_gpu(res, 0, index, co) -index_gpu_raft.add(xb) -D, I = index_gpu_raft.search(xq, 10) -D, I = index_gpu_raft.search(xq, 200) -print(D, I)