From 57d3b12501dc6382c6e7530a1b1f96ddd20e56d0 Mon Sep 17 00:00:00 2001 From: Mickael Ide Date: Mon, 16 Oct 2023 16:10:33 +0200 Subject: [PATCH] Fix IVF-Flat filtering --- .../raft/neighbors/detail/ivf_flat_interleaved_scan-inl.cuh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 ab2844c97c..3e0d7a4f7c 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 @@ -737,10 +737,11 @@ RAFT_KERNEL __launch_bounds__(kThreadsPerBlock) // This is the vector a given lane/thread handles const uint32_t vec_id = group_id * WarpSize + lane_id; - const bool valid = vec_id < list_length; + const bool valid = + vec_id < list_length && sample_filter(queries_offset + blockIdx.y, list_id, vec_id); // Process first shm_assisted_dim dimensions (always using shared memory) - if (valid && sample_filter(queries_offset + blockIdx.y, list_id, vec_id)) { + if (valid) { loadAndComputeDist lc(dist, compute_dist); for (int pos = 0; pos < shm_assisted_dim;