From d9742be1b2bbf1057afd2d3295f42c726960adf0 Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 23 Aug 2023 11:44:07 +0200 Subject: [PATCH 1/3] Fix the cluster probes (coarse_index) not being advanced when batching --- .../raft/neighbors/detail/ivf_flat_interleaved_scan-inl.cuh | 1 + 1 file changed, 1 insertion(+) 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; } } From 5d165c53cba6aeb7ac0b8f469dee2f79b571c494 Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 23 Aug 2023 11:45:47 +0200 Subject: [PATCH 2/3] Add the test cases --- cpp/test/neighbors/ann_ivf_flat.cuh | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh index d72d73680a..c4c8d727e1 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 From 728e55d1e292a43015eb67826d8887b15835aaf1 Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 23 Aug 2023 12:49:40 +0200 Subject: [PATCH 3/3] Fix a typo missing comma --- cpp/test/neighbors/ann_ivf_flat.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh index c4c8d727e1..71d48cdeb7 100644 --- a/cpp/test/neighbors/ann_ivf_flat.cuh +++ b/cpp/test/neighbors/ann_ivf_flat.cuh @@ -497,7 +497,7 @@ 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