File tree Expand file tree Collapse file tree 1 file changed +8
-7
lines changed
Expand file tree Collapse file tree 1 file changed +8
-7
lines changed Original file line number Diff line number Diff line change 11/*
2- * SPDX-FileCopyrightText: Copyright (c) 2023-2026 , NVIDIA CORPORATION.
2+ * SPDX-FileCopyrightText: Copyright (c) 2023-2025 , NVIDIA CORPORATION.
33 * SPDX-License-Identifier: Apache-2.0
44 */
55#pragma once
@@ -387,18 +387,19 @@ void bench_search(::benchmark::State& state,
387387 TODO: consider generating the filtered ground truth on-the-fly
388388 */
389389 uint32_t filter_pass_count = 0 ;
390- std::unordered_set<std::int32_t > gt_set;
391390 for (std::uint32_t l = 0 ; l < max_k && filter_pass_count < k; l++) {
392391 auto exp_idx = gt[i_orig_idx * max_k + l];
393392 if (!filter (exp_idx)) { continue ; }
394- gt_set.insert (exp_idx);
395393 filter_pass_count++;
394+ for (std::uint32_t j = 0 ; j < k; j++) {
395+ auto act_idx = static_cast <std::int32_t >(neighbors_host[i_out_idx * k + j]);
396+ if (act_idx == exp_idx) {
397+ match_count++;
398+ break ;
399+ }
400+ }
396401 }
397402 total_count += filter_pass_count;
398- for (std::uint32_t j = 0 ; j < k; j++) {
399- auto act_idx = static_cast <std::int32_t >(neighbors_host[i_out_idx * k + j]);
400- if (gt_set.count (act_idx)) match_count++;
401- }
402403 }
403404 }
404405 out_offset += n_queries;
You can’t perform that action at this time.
0 commit comments