From 3d35cdcb4c509a3f437f3b570b82c878882aedc2 Mon Sep 17 00:00:00 2001 From: Hiroyuki Ootomo Date: Sun, 24 Nov 2024 00:33:40 +0900 Subject: [PATCH 01/13] Fix CAGRA filter --- .../cagra/search_single_cta_kernel-inl.cuh | 104 +++++++++++++----- 1 file changed, 78 insertions(+), 26 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 79cb6bc101..aee6d6f6a1 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -453,6 +453,48 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort( MULTI_WARPS_2); } +// This function move the invalid index element to the end of the itopk list. +// Require : array_length % 32 == 0 && The invalid entry is only one. +template +RAFT_DEVICE_INLINE_FUNCTION void move_invalid_to_end_of_list(IdxT* const index_array, + float* const distance_array, + const std::uint32_t array_length) +{ + constexpr std::uint32_t warp_size = 32; + constexpr std::uint32_t invalid_index = utils::get_max_value(); + const std::uint32_t lane_id = threadIdx.x % warp_size; + + if (threadIdx.x >= warp_size) { return; } + + bool found_invalid = false; + if (array_length % warp_size == 0) { + for (std::uint32_t i = lane_id; i < array_length; i += warp_size) { + const auto index = index_array[i]; + const auto distance = distance_array[i]; + + if (found_invalid) { + index_array[i - 1] = index; + distance_array[i - 1] = distance; + } else { + // Check if the index is invalid + const auto I_found_invalid = (index == invalid_index); + const auto who_has_invalid = __ballot_sync(~0u, I_found_invalid); + // if a value that is loaded by a smaller lane id thread, shift the array + if (who_has_invalid << (warp_size - lane_id)) { + index_array[i - 1] = index; + distance_array[i - 1] = distance; + } + + found_invalid = who_has_invalid; + } + } + } + if (lane_id == 0) { + index_array[array_length - 1] = invalid_index; + distance_array[array_length - 1] = -1; + } +} + template RAFT_DEVICE_INLINE_FUNCTION void hashmap_restore(INDEX_T* const hashmap_ptr, const size_t hashmap_bitlen, @@ -627,28 +669,24 @@ __device__ void search_core( // topk with bitonic sort _CLK_START(); - if (std::is_same::value || - *filter_flag == 0) { - topk_by_bitonic_sort(result_distances_buffer, - result_indices_buffer, - internal_topk, - result_distances_buffer + internal_topk, - result_indices_buffer + internal_topk, - search_width * graph_degree, - topk_ws, - (iter == 0), - multi_warps_1, - multi_warps_2); - __syncthreads(); - } else { - topk_by_bitonic_sort_1st( - result_distances_buffer, - result_indices_buffer, - internal_topk + search_width * graph_degree, - internal_topk, - false); + if (!(std::is_same::value || + *filter_flag == 0)) { + // Move the filtered out index to the end of the itopk list + move_invalid_to_end_of_list(result_indices_buffer, result_distances_buffer, internal_topk); + if (threadIdx.x == 0) { *terminate_flag = 0; } } + topk_by_bitonic_sort(result_distances_buffer, + result_indices_buffer, + internal_topk, + result_distances_buffer + internal_topk, + result_indices_buffer + internal_topk, + search_width * graph_degree, + topk_ws, + (iter == 0), + multi_warps_1, + multi_warps_2); + __syncthreads(); _CLK_REC(clk_topk); } else { _CLK_START(); @@ -755,12 +793,26 @@ __device__ void search_core( } __syncthreads(); - topk_by_bitonic_sort_1st( - result_distances_buffer, - result_indices_buffer, - internal_topk + search_width * graph_degree, - top_k, - false); + topk_by_bitonic_sort_1st( + result_distances_buffer, result_indices_buffer, internal_topk, internal_topk, false); + + // If the sufficient number of valid indexes are not in the internal topk, pick up from the + // candidate list. + if (result_indices_buffer[top_k - 1] == invalid_index) { + __syncthreads(); + const unsigned multi_warps_1 = ((blockDim.x >= 64) && (MAX_CANDIDATES > 128)) ? 1 : 0; + const unsigned multi_warps_2 = ((blockDim.x >= 64) && (MAX_ITOPK > 256)) ? 1 : 0; + topk_by_bitonic_sort(result_distances_buffer, + result_indices_buffer, + internal_topk, + result_distances_buffer + internal_topk, + result_indices_buffer + internal_topk, + search_width * graph_degree, + topk_ws, + (iter == 0), + multi_warps_1, + multi_warps_2); + } __syncthreads(); } From ed10ed2c69880343f9cf93cafd22e15f077e9663 Mon Sep 17 00:00:00 2001 From: Hiroyuki Ootomo Date: Sun, 24 Nov 2024 00:56:53 +0900 Subject: [PATCH 02/13] Fix new distance --- cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index aee6d6f6a1..53c2b0bc91 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -491,7 +491,7 @@ RAFT_DEVICE_INLINE_FUNCTION void move_invalid_to_end_of_list(IdxT* const index_a } if (lane_id == 0) { index_array[array_length - 1] = invalid_index; - distance_array[array_length - 1] = -1; + distance_array[array_length - 1] = utils::get_max_value(); } } From f476fd7a5c80d02cbd48d49cccc8a11fae039f58 Mon Sep 17 00:00:00 2001 From: Hiroyuki Ootomo Date: Tue, 26 Nov 2024 11:39:42 +0900 Subject: [PATCH 03/13] Fix CAGRA filtering test --- cpp/test/neighbors/ann_cagra.cuh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/cpp/test/neighbors/ann_cagra.cuh b/cpp/test/neighbors/ann_cagra.cuh index 660246c67e..82db66914a 100644 --- a/cpp/test/neighbors/ann_cagra.cuh +++ b/cpp/test/neighbors/ann_cagra.cuh @@ -758,11 +758,10 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { search_params.algo = ps.algo; search_params.max_queries = ps.max_queries; search_params.team_size = ps.team_size; + search_params.itopk_size = ps.itopk_size; - // TODO: setting search_params.itopk_size here breaks the filter tests, but is required for // k>1024 skip these tests until fixed if (ps.k >= 1024) { GTEST_SKIP(); } - // search_params.itopk_size = ps.itopk_size; auto database_view = raft::make_device_matrix_view( (const DataT*)database.data(), ps.n_rows, ps.dim); From 692375d46026d9a1df73cb2829c96c370931d374 Mon Sep 17 00:00:00 2001 From: Hiroyuki Ootomo Date: Tue, 26 Nov 2024 15:56:53 +0900 Subject: [PATCH 04/13] Make thew CAGRA filtering postprocess sort-free --- .../cagra/search_single_cta_kernel-inl.cuh | 36 +++++++++++++++++-- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 53c2b0bc91..16a5f1473e 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -793,12 +793,42 @@ __device__ void search_core( } __syncthreads(); - topk_by_bitonic_sort_1st( - result_distances_buffer, result_indices_buffer, internal_topk, internal_topk, false); + // Move invalid index items to the end of the buffer without sorting the entire buffer + using scan_op_t = cub::WarpScan; + auto& temp_storage = *reinterpret_cast(smem_work_ptr); + + constexpr std::uint32_t warp_size = 32; + if (threadIdx.x < warp_size) { + std::uint32_t position_offset = 0; + for (std::uint32_t buffer_offset = 0; buffer_offset < internal_topk; + buffer_offset += warp_size) { + // Calculate the new buffer index + const auto src_position = buffer_offset + threadIdx.x; + const std::uint32_t is_valid_index = + (result_indices_buffer[src_position] & (~index_msb_1_mask)) == invalid_index ? 0 : 1; + std::uint32_t new_position; + scan_op_t(temp_storage).InclusiveSum(is_valid_index, new_position); + if (is_valid_index) { + const auto dst_position = position_offset + (new_position - 1); + result_indices_buffer[dst_position] = result_indices_buffer[src_position]; + result_distances_buffer[dst_position] = result_distances_buffer[src_position]; + } + + // Calculate the largest valid position within a warp and bcast it for the next iteration + position_offset += new_position; + for (std::uint32_t offset = (warp_size >> 1); offset > 0; offset >>= 1) { + const auto v = __shfl_xor_sync(~0u, position_offset, offset); + if ((threadIdx.x & offset) == 0) { position_offset = v; } + } + + // If the enough number of items are found, do early termination + if (position_offset >= top_k) { break; } + } + } // If the sufficient number of valid indexes are not in the internal topk, pick up from the // candidate list. - if (result_indices_buffer[top_k - 1] == invalid_index) { + if (top_k > internal_topk || result_indices_buffer[top_k - 1] == invalid_index) { __syncthreads(); const unsigned multi_warps_1 = ((blockDim.x >= 64) && (MAX_CANDIDATES > 128)) ? 1 : 0; const unsigned multi_warps_2 = ((blockDim.x >= 64) && (MAX_ITOPK > 256)) ? 1 : 0; From 67169400dba8e5ac8f6f09c840a7c5141a8de7db Mon Sep 17 00:00:00 2001 From: Hiroyuki Ootomo Date: Tue, 26 Nov 2024 16:17:50 +0900 Subject: [PATCH 05/13] Update CAGRA filtering postprocess --- .../cagra/search_single_cta_kernel-inl.cuh | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 16a5f1473e..23d6aae06e 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -799,7 +799,7 @@ __device__ void search_core( constexpr std::uint32_t warp_size = 32; if (threadIdx.x < warp_size) { - std::uint32_t position_offset = 0; + std::uint32_t num_found_valid = 0; for (std::uint32_t buffer_offset = 0; buffer_offset < internal_topk; buffer_offset += warp_size) { // Calculate the new buffer index @@ -809,20 +809,29 @@ __device__ void search_core( std::uint32_t new_position; scan_op_t(temp_storage).InclusiveSum(is_valid_index, new_position); if (is_valid_index) { - const auto dst_position = position_offset + (new_position - 1); + const auto dst_position = num_found_valid + (new_position - 1); result_indices_buffer[dst_position] = result_indices_buffer[src_position]; result_distances_buffer[dst_position] = result_distances_buffer[src_position]; } // Calculate the largest valid position within a warp and bcast it for the next iteration - position_offset += new_position; + num_found_valid += new_position; for (std::uint32_t offset = (warp_size >> 1); offset > 0; offset >>= 1) { - const auto v = __shfl_xor_sync(~0u, position_offset, offset); - if ((threadIdx.x & offset) == 0) { position_offset = v; } + const auto v = __shfl_xor_sync(~0u, num_found_valid, offset); + if ((threadIdx.x & offset) == 0) { num_found_valid = v; } } // If the enough number of items are found, do early termination - if (position_offset >= top_k) { break; } + if (num_found_valid >= top_k) { break; } + } + + if (num_found_valid < top_k) { + // Fill the remaining buffer with invalid values so that `topk_by_bitonic_sort` is usable in + // the next step + for (std::uint32_t i = num_found_valid + threadIdx.x; i < internal_topk; i += warp_size) { + result_indices_buffer[i] = invalid_index; + result_distances_buffer[i] = utils::get_max_value(); + } } } From 6c75170cf6dc6b0909346843d5d71bb0f7beba33 Mon Sep 17 00:00:00 2001 From: Hiroyuki Ootomo Date: Tue, 26 Nov 2024 20:42:35 +0900 Subject: [PATCH 06/13] Fix test --- cpp/test/neighbors/ann_cagra.cuh | 3 --- 1 file changed, 3 deletions(-) diff --git a/cpp/test/neighbors/ann_cagra.cuh b/cpp/test/neighbors/ann_cagra.cuh index 82db66914a..8d57014395 100644 --- a/cpp/test/neighbors/ann_cagra.cuh +++ b/cpp/test/neighbors/ann_cagra.cuh @@ -760,9 +760,6 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { search_params.team_size = ps.team_size; search_params.itopk_size = ps.itopk_size; - // k>1024 skip these tests until fixed - if (ps.k >= 1024) { GTEST_SKIP(); } - auto database_view = raft::make_device_matrix_view( (const DataT*)database.data(), ps.n_rows, ps.dim); From b87b0279f44ab73da1649088e2a1b4ea6a5e47ec Mon Sep 17 00:00:00 2001 From: Hiroyuki Ootomo Date: Tue, 26 Nov 2024 20:43:40 +0900 Subject: [PATCH 07/13] Fix invalid node move --- .../neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 23d6aae06e..62086e3982 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -672,7 +672,10 @@ __device__ void search_core( if (!(std::is_same::value || *filter_flag == 0)) { // Move the filtered out index to the end of the itopk list - move_invalid_to_end_of_list(result_indices_buffer, result_distances_buffer, internal_topk); + for (unsigned i = 0; i < search_width; i++) { + move_invalid_to_end_of_list( + result_indices_buffer, result_distances_buffer, internal_topk); + } if (threadIdx.x == 0) { *terminate_flag = 0; } } From 4ad78b378b6e992427099197c93a99da5a632741 Mon Sep 17 00:00:00 2001 From: tsuki <12711693+enp1s0@users.noreply.github.com> Date: Tue, 3 Dec 2024 23:27:19 +0900 Subject: [PATCH 08/13] Update cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh Co-authored-by: Artem M. Chirkin <9253178+achirkin@users.noreply.github.com> --- cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 62086e3982..f63231f177 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -478,7 +478,7 @@ RAFT_DEVICE_INLINE_FUNCTION void move_invalid_to_end_of_list(IdxT* const index_a } else { // Check if the index is invalid const auto I_found_invalid = (index == invalid_index); - const auto who_has_invalid = __ballot_sync(~0u, I_found_invalid); + const auto who_has_invalid = raft::ballot(I_found_invalid); // if a value that is loaded by a smaller lane id thread, shift the array if (who_has_invalid << (warp_size - lane_id)) { index_array[i - 1] = index; From 536ed1f5cba2c6658636829a25044acc00398832 Mon Sep 17 00:00:00 2001 From: tsuki <12711693+enp1s0@users.noreply.github.com> Date: Tue, 3 Dec 2024 23:27:27 +0900 Subject: [PATCH 09/13] Update cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh Co-authored-by: Artem M. Chirkin <9253178+achirkin@users.noreply.github.com> --- cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index f63231f177..9f551844c6 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -820,7 +820,7 @@ __device__ void search_core( // Calculate the largest valid position within a warp and bcast it for the next iteration num_found_valid += new_position; for (std::uint32_t offset = (warp_size >> 1); offset > 0; offset >>= 1) { - const auto v = __shfl_xor_sync(~0u, num_found_valid, offset); + const auto v = raft::shfl_xor( num_found_valid, offset); if ((threadIdx.x & offset) == 0) { num_found_valid = v; } } From 270b9285cd889283880a376479344c74715c8dc5 Mon Sep 17 00:00:00 2001 From: Hiroyuki Ootomo Date: Wed, 4 Dec 2024 00:05:54 +0900 Subject: [PATCH 10/13] Fix style --- cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 9f551844c6..cec6bdbb99 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -820,7 +820,7 @@ __device__ void search_core( // Calculate the largest valid position within a warp and bcast it for the next iteration num_found_valid += new_position; for (std::uint32_t offset = (warp_size >> 1); offset > 0; offset >>= 1) { - const auto v = raft::shfl_xor( num_found_valid, offset); + const auto v = raft::shfl_xor(num_found_valid, offset); if ((threadIdx.x & offset) == 0) { num_found_valid = v; } } From 3d676e2a45e14bfa5686cb1c92f560ed34ec28cc Mon Sep 17 00:00:00 2001 From: Hiroyuki Ootomo Date: Wed, 4 Dec 2024 00:27:00 +0900 Subject: [PATCH 11/13] Fix smem size calculation for filtering postprocess in single-CAT kernel --- .../neighbors/detail/cagra/search_single_cta.cuh | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh index 2bed190090..fa71dbaf95 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh @@ -129,17 +129,27 @@ struct search : search_plan_impl { (sizeof(INDEX_T) + sizeof(DISTANCE_T)) * result_buffer_size_32 + sizeof(INDEX_T) * hashmap::get_size(small_hash_bitlen) + sizeof(INDEX_T) * search_width + sizeof(std::uint32_t) * topk_ws_size + sizeof(std::uint32_t); - smem_size = base_smem_size; + + std::uint32_t additional_smem_size = 0; if (num_itopk_candidates > 256) { // Tentatively calculate the required share memory size when radix // sort based topk is used, assuming the block size is the maximum. if (itopk_size <= 256) { - smem_size += topk_by_radix_sort<256, INDEX_T>::smem_size * sizeof(std::uint32_t); + additional_smem_size += topk_by_radix_sort<256, INDEX_T>::smem_size * sizeof(std::uint32_t); } else { - smem_size += topk_by_radix_sort<512, INDEX_T>::smem_size * sizeof(std::uint32_t); + additional_smem_size += topk_by_radix_sort<512, INDEX_T>::smem_size * sizeof(std::uint32_t); } } + if (!std::is_same_v) { + // For filtering postprocess + using scan_op_t = cub::WarpScan; + additional_smem_size = + std::max(additional_smem_size, sizeof(scan_op_t::TempStorage)); + } + + smem_size = base_smem_size + additional_smem_size; + uint32_t block_size = thread_block_size; if (block_size == 0) { block_size = min_block_size; From 49acbabbd0414d4831d44b11c422ed69dc129b51 Mon Sep 17 00:00:00 2001 From: Hiroyuki Ootomo Date: Wed, 4 Dec 2024 00:38:21 +0900 Subject: [PATCH 12/13] Fix the topk function name --- .../cagra/search_single_cta_kernel-inl.cuh | 80 ++++++++++--------- 1 file changed, 41 insertions(+), 39 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index cec6bdbb99..d13b6e394a 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -111,7 +111,7 @@ RAFT_DEVICE_INLINE_FUNCTION void pickup_next_parents(std::uint32_t* const termin } template -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_1st( +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_full( float* candidate_distances, // [num_candidates] IdxT* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, @@ -215,7 +215,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_1st( } template -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_2nd( +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_merge( float* itopk_distances, // [num_itopk] IdxT* itopk_indices, // [num_itopk] const std::uint32_t num_itopk, @@ -424,7 +424,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_2nd( template -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort( +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( float* itopk_distances, // [num_itopk] IdxT* itopk_indices, // [num_itopk] const std::uint32_t num_itopk, @@ -437,20 +437,20 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort( const unsigned MULTI_WARPS_2) { // The results in candidate_distances/indices are sorted by bitonic sort. - topk_by_bitonic_sort_1st( + topk_by_bitonic_sort_and_merge_full( candidate_distances, candidate_indices, num_candidates, num_itopk, MULTI_WARPS_1); // The results sorted above are merged with the internal intermediate top-k // results so far using bitonic merge. - topk_by_bitonic_sort_2nd(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - MULTI_WARPS_2); + topk_by_bitonic_sort_and_merge_merge(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + MULTI_WARPS_2); } // This function move the invalid index element to the end of the itopk list. @@ -631,10 +631,10 @@ __device__ void search_core( // sort if constexpr (TOPK_BY_BITONIC_SORT) { // [Notice] - // It is good to use multiple warps in topk_by_bitonic_sort() when + // It is good to use multiple warps in topk_by_bitonic_sort_and_merge() when // batch size is small (short-latency), but it might not be always good // when batch size is large (high-throughput). - // topk_by_bitonic_sort() consists of two operations: + // topk_by_bitonic_sort_and_merge() consists of two operations: // if MAX_CANDIDATES is greater than 128, the first operation uses two warps; // if MAX_ITOPK is greater than 256, the second operation used two warps. const unsigned multi_warps_1 = ((blockDim.x >= 64) && (MAX_CANDIDATES > 128)) ? 1 : 0; @@ -643,9 +643,9 @@ __device__ void search_core( // reset small-hash table. if ((iter + 1) % small_hash_reset_interval == 0) { // Depending on the block size and the number of warps used in - // topk_by_bitonic_sort(), determine which warps are used to reset + // topk_by_bitonic_sort_and_merge(), determine which warps are used to reset // the small hash and whether they are performed in overlap with - // topk_by_bitonic_sort(). + // topk_by_bitonic_sort_and_merge(). _CLK_START(); unsigned hash_start_tid; if (blockDim.x == 32) { @@ -679,16 +679,17 @@ __device__ void search_core( if (threadIdx.x == 0) { *terminate_flag = 0; } } - topk_by_bitonic_sort(result_distances_buffer, - result_indices_buffer, - internal_topk, - result_distances_buffer + internal_topk, - result_indices_buffer + internal_topk, - search_width * graph_degree, - topk_ws, - (iter == 0), - multi_warps_1, - multi_warps_2); + topk_by_bitonic_sort_and_merge( + result_distances_buffer, + result_indices_buffer, + internal_topk, + result_distances_buffer + internal_topk, + result_indices_buffer + internal_topk, + search_width * graph_degree, + topk_ws, + (iter == 0), + multi_warps_1, + multi_warps_2); __syncthreads(); _CLK_REC(clk_topk); } else { @@ -829,8 +830,8 @@ __device__ void search_core( } if (num_found_valid < top_k) { - // Fill the remaining buffer with invalid values so that `topk_by_bitonic_sort` is usable in - // the next step + // Fill the remaining buffer with invalid values so that `topk_by_bitonic_sort_and_merge` is + // usable in the next step for (std::uint32_t i = num_found_valid + threadIdx.x; i < internal_topk; i += warp_size) { result_indices_buffer[i] = invalid_index; result_distances_buffer[i] = utils::get_max_value(); @@ -844,16 +845,17 @@ __device__ void search_core( __syncthreads(); const unsigned multi_warps_1 = ((blockDim.x >= 64) && (MAX_CANDIDATES > 128)) ? 1 : 0; const unsigned multi_warps_2 = ((blockDim.x >= 64) && (MAX_ITOPK > 256)) ? 1 : 0; - topk_by_bitonic_sort(result_distances_buffer, - result_indices_buffer, - internal_topk, - result_distances_buffer + internal_topk, - result_indices_buffer + internal_topk, - search_width * graph_degree, - topk_ws, - (iter == 0), - multi_warps_1, - multi_warps_2); + topk_by_bitonic_sort_and_merge( + result_distances_buffer, + result_indices_buffer, + internal_topk, + result_distances_buffer + internal_topk, + result_indices_buffer + internal_topk, + search_width * graph_degree, + topk_ws, + (iter == 0), + multi_warps_1, + multi_warps_2); } __syncthreads(); } From 55afa103de55de7fef981de993c8826f5b7aea9a Mon Sep 17 00:00:00 2001 From: Hiroyuki Ootomo Date: Wed, 4 Dec 2024 00:40:47 +0900 Subject: [PATCH 13/13] Fix the topk function name (2) --- .../cagra/search_single_cta_kernel-inl.cuh | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index d13b6e394a..678ed0cb45 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -111,7 +111,7 @@ RAFT_DEVICE_INLINE_FUNCTION void pickup_next_parents(std::uint32_t* const termin } template -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_full( +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( float* candidate_distances, // [num_candidates] IdxT* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, @@ -215,7 +215,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_full( } template -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_merge( +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( float* itopk_distances, // [num_itopk] IdxT* itopk_indices, // [num_itopk] const std::uint32_t num_itopk, @@ -437,20 +437,20 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( const unsigned MULTI_WARPS_2) { // The results in candidate_distances/indices are sorted by bitonic sort. - topk_by_bitonic_sort_and_merge_full( + topk_by_bitonic_sort_and_full( candidate_distances, candidate_indices, num_candidates, num_itopk, MULTI_WARPS_1); // The results sorted above are merged with the internal intermediate top-k // results so far using bitonic merge. - topk_by_bitonic_sort_and_merge_merge(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - MULTI_WARPS_2); + topk_by_bitonic_sort_and_merge(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + MULTI_WARPS_2); } // This function move the invalid index element to the end of the itopk list.