From a782cd0201f5d1349b426a87b0493cf140298325 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 17:12:45 +0000 Subject: [PATCH] Handle empty lists when building IVF-RaBitQ indexes k-means can leave lists empty, for example when the dataset has fewer distinct rows than n_lists. Building from device data, or from host data that fits on the device, then launched the quantization kernels with zero blocks and failed with cudaErrorInvalidValue. The streaming build skipped empty lists, which left their rotated centroids uninitialized. Compute only the rotated centroid for an empty list and skip the kernels that have no rows to process. Add a regression test with fewer distinct rows than lists. --- .../neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu | 54 ++++---- .../ivf_rabitq/gpu_index/quantizer_gpu.cu | 12 ++ cpp/tests/neighbors/ann_ivf_rabitq.cuh | 119 ++++++++++++++++++ .../ann_ivf_rabitq/test_float_int64_t.cu | 5 + 4 files changed, 166 insertions(+), 24 deletions(-) diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu index 9ce10fd57e..70371ddc54 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -719,38 +719,44 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, cluster_idx++; } - if (batch_vectors == 0) break; // No more clusters to process + // Every cluster fits into a batch (checked above), so each batch makes progress. + RAFT_EXPECTS(cluster_idx > batch_start_cluster, "IVF-RaBitQ streaming batch is empty"); // ------------------------- // 11. Gather and transfer batch data to GPU using reordered PIDs // ------------------------- - // First, copy the reordered PIDs for this batch to host (reuse preallocated buffer) - raft::copy(batch_pids.data_handle(), d_flat_pids + batch_start_offset, batch_vectors, stream_); - raft::resource::sync_stream(handle_); + // A batch that only contains empty clusters has no data to transfer, but the clusters still + // need their rotated centroids below. + if (batch_vectors > 0) { + // First, copy the reordered PIDs for this batch to host (reuse preallocated buffer) + raft::copy( + batch_pids.data_handle(), d_flat_pids + batch_start_offset, batch_vectors, stream_); + raft::resource::sync_stream(handle_); - // Gather data on host using the reordered PIDs (reuse preallocated buffer) - // Use OpenMP parallel region to leverage persistent thread pool across batches + // Gather data on host using the reordered PIDs (reuse preallocated buffer) + // Use OpenMP parallel region to leverage persistent thread pool across batches #pragma omp parallel num_threads(num_threads) - { - size_t tid = cuvs::core::omp::get_thread_num(); - size_t chunk_size = batch_vectors / num_threads; - size_t start = tid * chunk_size; - size_t end = (tid == num_threads - 1) ? batch_vectors : start + chunk_size; - - for (size_t i = start; i < end; ++i) { - std::memcpy(&batch_data.data_handle()[i * num_dimensions], - &host_data[batch_pids.data_handle()[i] * num_dimensions], - num_dimensions * sizeof(float)); + { + size_t tid = cuvs::core::omp::get_thread_num(); + size_t chunk_size = batch_vectors / num_threads; + size_t start = tid * chunk_size; + size_t end = (tid == num_threads - 1) ? batch_vectors : start + chunk_size; + + for (size_t i = start; i < end; ++i) { + std::memcpy(&batch_data.data_handle()[i * num_dimensions], + &host_data[batch_pids.data_handle()[i] * num_dimensions], + num_dimensions * sizeof(float)); + } } - } - // OpenMP implicit barrier ensures all threads complete before continuing + // OpenMP implicit barrier ensures all threads complete before continuing - // Transfer batch to GPU (reuse preallocated buffer) - raft::copy(d_batch_data.data_handle(), - batch_data.data_handle(), - batch_vectors * num_dimensions, - stream_); + // Transfer batch to GPU (reuse preallocated buffer) + raft::copy(d_batch_data.data_handle(), + batch_data.data_handle(), + batch_vectors * num_dimensions, + stream_); + } // ------------------------- // 12. Quantize each cluster in the batch @@ -758,8 +764,8 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, size_t batch_offset = 0; for (size_t c = batch_start_cluster; c < cluster_idx; ++c) { + // Empty clusters are quantized too, so that their rotated centroids are initialized. size_t cluster_size = h_cluster_meta[c].num; - if (cluster_size == 0) continue; const float* cur_centroid = device_centroids + c * num_dimensions; float* cur_rotated_c = d_rotated_centroids.data_handle() + c * num_padded_dim; diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu index b35ac9fab9..5eef2179d4 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu @@ -520,6 +520,9 @@ void data_transformation_batch_opt(const float* d_data, // 5. Save the rotated centroid: copy CP into d_rotated_c. raft::copy(d_rotated_c, d_CP, D, stream); + // An empty list only needs its rotated centroid; a launch with zero blocks is invalid. + if (num_points == 0) { return; } + // 6. Launch the single FUSED kernel for subtract, normalize, and binarize. const unsigned int FusedBlockSize = 256; // A good default, can be tuned. dim3 gridDim(num_points); @@ -637,6 +640,9 @@ void DataQuantizerGPU::quantize_batch_opt(const float* d_data, D, handle_); + // An empty list has no codes or factors to compute. + if (num_points == 0) { return; } + rabitq_codes_and_factors_fused(d_rotated_c, d_bin_XP.data_handle(), d_XP.data_handle(), @@ -714,6 +720,9 @@ void data_transformation_batch_opt_contiguous(const float* d_contiguous_data, // 5. Save the rotated centroid: copy CP into d_rotated_c. raft::copy(d_rotated_c, d_CP, D, stream); + // An empty list only needs its rotated centroid; a launch with zero blocks is invalid. + if (num_points == 0) { return; } + // 6. Launch the single FUSED kernel for subtract, normalize, and binarize. const unsigned int FusedBlockSize = 256; // A good default, can be tuned. dim3 gridDim(num_points); @@ -758,6 +767,9 @@ void DataQuantizerGPU::quantize_batch_opt_contiguous(const float* d_contiguous_d D, handle_); + // An empty list has no codes or factors to compute. + if (num_points == 0) { return; } + rabitq_codes_and_factors_fused(d_rotated_c, d_bin_XP.data_handle(), d_XP.data_handle(), diff --git a/cpp/tests/neighbors/ann_ivf_rabitq.cuh b/cpp/tests/neighbors/ann_ivf_rabitq.cuh index 2755258600..7296e67ddf 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq.cuh +++ b/cpp/tests/neighbors/ann_ivf_rabitq.cuh @@ -396,4 +396,123 @@ inline auto var_search_mode_1_bit() -> test_cases_t #define INSTANTIATE(type, vals) \ INSTANTIATE_TEST_SUITE_P(IvfRabitq, type, ::testing::ValuesIn(vals)); /* NOLINT */ +/* Empty lists */ + +enum class dataset_input { device, host, host_streaming }; + +struct ivf_rabitq_empty_lists_inputs { + uint32_t dim; + dataset_input input; +}; + +inline auto operator<<(std::ostream& os, const ivf_rabitq_empty_lists_inputs& p) -> std::ostream& +{ + const char* input = p.input == dataset_input::device ? "device" + : p.input == dataset_input::host ? "host" + : "host_streaming"; + return os << "ivf_rabitq_empty_lists_inputs {.dim = " << p.dim << ", .input = " << input << "}"; +} + +/** + * The dataset has fewer distinct rows than there are lists. Identical rows are always assigned to + * the same list, so k-means is guaranteed to leave some lists empty. The index must still build + * from device and host input, and search must still find the nearest rows. + */ +class ivf_rabitq_empty_lists_test : public ::testing::TestWithParam { + protected: + void run() + { + constexpr int64_t kRows = 4096; + constexpr int64_t kDistinct = 8; + constexpr int64_t kQueries = 2 * kDistinct; + constexpr int64_t kK = 10; + constexpr uint32_t kLists = 32; + constexpr float kSpacing = 0.25f; + constexpr float kOffset = 0.1f; + + auto ps = GetParam(); + auto dim = static_cast(ps.dim); + auto stream = raft::resource::get_cuda_stream(handle_); + + // Row i is the point kSpacing * (i % kDistinct) on the diagonal. Query q is offset from the + // distinct row q % kDistinct by +-kOffset in every dimension, so that row is its unique nearest + // distinct row at squared L2 distance kOffset^2 * dim. + auto h_dataset = raft::make_host_matrix(kRows, dim); + for (int64_t i = 0; i < kRows; i++) { + for (int64_t j = 0; j < dim; j++) { + h_dataset(i, j) = kSpacing * static_cast(i % kDistinct); + } + } + auto h_queries = raft::make_host_matrix(kQueries, dim); + for (int64_t i = 0; i < kQueries; i++) { + for (int64_t j = 0; j < dim; j++) { + h_queries(i, j) = + kSpacing * static_cast(i % kDistinct) + (i < kDistinct ? kOffset : -kOffset); + } + } + auto d_dataset = raft::make_device_matrix(handle_, kRows, dim); + auto d_queries = raft::make_device_matrix(handle_, kQueries, dim); + raft::copy(d_dataset.data_handle(), h_dataset.data_handle(), h_dataset.size(), stream); + raft::copy(d_queries.data_handle(), h_queries.data_handle(), h_queries.size(), stream); + + index_params ipams; + ipams.n_lists = kLists; + if (ps.input == dataset_input::host_streaming) { + ipams.force_streaming = true; + ipams.streaming_batch_size = kRows; + } + auto idx = ps.input == dataset_input::device + ? build(handle_, ipams, raft::make_const_mdspan(d_dataset.view())) + : build(handle_, ipams, raft::make_const_mdspan(h_dataset.view())); + + // Probe every list, because an empty list's centroid can be closer to a query than the + // centroid of the list that holds its neighbors. + search_params sparams; + sparams.n_probes = kLists; + auto neighbors = raft::make_device_matrix(handle_, kQueries, kK); + auto distances = raft::make_device_matrix(handle_, kQueries, kK); + search(handle_, + sparams, + idx, + raft::make_const_mdspan(d_queries.view()), + neighbors.view(), + distances.view()); + + std::vector h_neighbors(kQueries * kK); + std::vector h_distances(kQueries * kK); + raft::copy(h_neighbors.data(), neighbors.data_handle(), h_neighbors.size(), stream); + raft::copy(h_distances.data(), distances.data_handle(), h_distances.size(), stream); + raft::resource::sync_stream(handle_); + + // Every row has kRows / kDistinct > kK copies, so all kK neighbors are copies of the nearest + // distinct row. The next nearest distinct row is (kSpacing - kOffset)^2 * dim away. + const float expected_dist = kOffset * kOffset * static_cast(dim); + for (int64_t q = 0; q < kQueries; q++) { + for (int64_t j = 0; j < kK; j++) { + auto neighbor = h_neighbors[q * kK + j]; + auto dist = h_distances[q * kK + j]; + ASSERT_TRUE(neighbor >= 0 && neighbor < kRows) + << ps << ": query " << q << ", neighbor " << j << " = " << neighbor; + ASSERT_EQ(neighbor % kDistinct, q % kDistinct) + << ps << ": query " << q << ", neighbor " << j << " = " << neighbor; + ASSERT_NEAR(dist, expected_dist, 0.1f * expected_dist) + << ps << ": query " << q << ", neighbor " << j << " = " << neighbor; + } + } + } + + raft::resources handle_; +}; + +inline auto empty_lists_inputs() -> std::vector +{ + std::vector xs; + for (uint32_t dim : {1, 2, 7, 64, 100}) { + for (auto input : {dataset_input::device, dataset_input::host, dataset_input::host_streaming}) { + xs.push_back({dim, input}); + } + } + return xs; +} + } // namespace cuvs::neighbors::ivf_rabitq diff --git a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu index 42bf84e6a3..28771bd785 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu @@ -18,4 +18,9 @@ INSTANTIATE(f32_f32_i64, defaults() + small_dims() + big_dims() + var_n_probes() + var_k() + var_bits_per_dim() + var_search_mode() + var_search_mode_1_bit()); +TEST_P(ivf_rabitq_empty_lists_test, build_search) { this->run(); } // NOLINT +INSTANTIATE_TEST_SUITE_P(IvfRabitq, + ivf_rabitq_empty_lists_test, + ::testing::ValuesIn(empty_lists_inputs())); // NOLINT + } // namespace cuvs::neighbors::ivf_rabitq