From 9019054cd41f36639cb9c9552d91b6249ef7bab8 Mon Sep 17 00:00:00 2001 From: "Artem M. Chirkin" <9253178+achirkin@users.noreply.github.com> Date: Wed, 23 Aug 2023 15:15:49 +0200 Subject: [PATCH] IVF-Flat: fix search batching (#1764) Fix the cluster probes (coarse_index) not being advanced when batching. Thanks @tarang-jain for the precise reproducer. Closes: #1756 Authors: - Artem M. Chirkin (https://github.com/achirkin) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/raft/pull/1764 --- .../neighbors/detail/ivf_flat_interleaved_scan-inl.cuh | 1 + cpp/test/neighbors/ann_ivf_flat.cuh | 7 ++++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_interleaved_scan-inl.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_interleaved_scan-inl.cuh index 18f1906dc5..81779668c4 100644 --- a/cpp/include/raft/neighbors/detail/ivf_flat_interleaved_scan-inl.cuh +++ b/cpp/include/raft/neighbors/detail/ivf_flat_interleaved_scan-inl.cuh @@ -884,6 +884,7 @@ void launch_kernel(Lambda lambda, queries += grid_dim_y * index.dim(); neighbors += grid_dim_y * grid_dim_x * k; distances += grid_dim_y * grid_dim_x * k; + coarse_index += grid_dim_y * n_probes; } } diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh index d72d73680a..71d48cdeb7 100644 --- a/cpp/test/neighbors/ann_ivf_flat.cuh +++ b/cpp/test/neighbors/ann_ivf_flat.cuh @@ -497,6 +497,11 @@ const std::vector> inputs = { raft::matrix::detail::select::warpsort::kMaxCapacity * 4, raft::matrix::detail::select::warpsort::kMaxCapacity * 4, raft::distance::DistanceType::InnerProduct, - false}}; + false}, + + // The following two test cases should show very similar recall. + // num_queries, num_db_vecs, dim, k, nprobe, nlist, metric, adaptive_centers + {20000, 8712, 3, 10, 51, 66, raft::distance::DistanceType::L2Expanded, false}, + {100000, 8712, 3, 10, 51, 66, raft::distance::DistanceType::L2Expanded, false}}; } // namespace raft::neighbors::ivf_flat