From 6aa64e866d5299a1038ac1c245d75f3f665dea99 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Sat, 21 Sep 2024 18:54:51 -0700 Subject: [PATCH 01/24] all changes --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 25 ++++++++++++++----- .../ivf_pq/ivf_pq_compute_similarity_impl.cuh | 2 ++ cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 11 +++++++- cpp/test/neighbors/ann_ivf_pq.cuh | 20 +++++++++++++++ .../ann_ivf_pq/test_float_int64_t.cu | 4 +-- 5 files changed, 53 insertions(+), 9 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index c65ea8108..09e1396b3 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -41,6 +41,8 @@ #include #include #include +#include +#include #include #include #include @@ -1569,17 +1571,22 @@ void extend(raft::resources const& handle, stream)); vec_batches.prefetch_next_batch(); for (const auto& batch : vec_batches) { - auto batch_data_view = raft::make_device_matrix_view( - batch.data(), batch.size(), index->dim()); + // auto batch_data_view = raft::make_device_matrix_view( + // batch.data(), batch.size(), index->dim()); auto batch_labels_view = raft::make_device_vector_view( new_data_labels.data() + batch.offset(), batch.size()); + auto float_vec_batch = raft::make_device_matrix(handle, batch.size(), index->dim()); + raft::linalg::map_offset(handle, raft::make_device_vector_view(batch.data(), batch.size() * index->dim()), raft::make_device_vector_view(float_vec_batch.data_handle(), float_vec_batch.size()), [=]__device__(internal_extents_t idx, T i) {return utils::mapping{}(i);}); + if(index->metric() == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_vec_batch.view()), float_vec_batch.view(), raft::linalg::NormType::L2Norm); + } auto centers_view = raft::make_device_matrix_view( cluster_centers.data(), n_clusters, index->dim()); cuvs::cluster::kmeans::balanced_params kmeans_params; - kmeans_params.metric = static_cast((int)index->metric()); + kmeans_params.metric = cuvs::distance::DistanceType::InnerProduct; cuvs::cluster::kmeans_balanced::predict(handle, kmeans_params, - batch_data_view, + raft::make_const_mdspan(float_vec_batch.view()), centers_view, batch_labels_view, utils::mapping{}); @@ -1632,9 +1639,14 @@ void extend(raft::resources const& handle, vec_batches.prefetch_next_batch(); for (const auto& vec_batch : vec_batches) { const auto& idx_batch = *idx_batches++; + auto float_vec_batch = raft::make_device_matrix(handle, vec_batch.size(), index->dim()); + raft::linalg::map_offset(handle, raft::make_device_vector_view(vec_batch.data(), vec_batch.size() * index->dim()), raft::make_device_vector_view(float_vec_batch.data_handle(), float_vec_batch.size()), [=]__device__(internal_extents_t idx, T i) {return utils::mapping{}(i);}); + if(index->metric() == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_vec_batch.view()), float_vec_batch.view(), raft::linalg::NormType::L2Norm); + } process_and_fill_codes(handle, *index, - vec_batch.data(), + float_vec_batch.data_handle(), new_indices != nullptr ? std::variant(idx_batch.data()) : std::variant(IdxT(idx_batch.offset())), @@ -1750,11 +1762,12 @@ auto build(raft::resources const& handle, // Train balanced hierarchical kmeans clustering auto trainset_const_view = raft::make_const_mdspan(trainset.view()); + raft::linalg::row_normalize(handle, trainset_const_view, trainset.view(), raft::linalg::NormType::L2Norm); auto centers_view = raft::make_device_matrix_view( cluster_centers, index.n_lists(), index.dim()); cuvs::cluster::kmeans::balanced_params kmeans_params; kmeans_params.n_iters = params.kmeans_n_iters; - kmeans_params.metric = static_cast((int)index.metric()); + kmeans_params.metric = cuvs::distance::DistanceType::InnerProduct; cuvs::cluster::kmeans_balanced::fit( handle, kmeans_params, trainset_const_view, centers_view, utils::mapping{}); diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh index 5fccbb385..7a5a2336b 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh @@ -369,6 +369,7 @@ RAFT_KERNEL compute_similarity_kernel(uint32_t dim, reinterpret_cast(lut_end)[i] = query[i] - cluster_center[i]; } } break; + case distance::DistanceType::CosineExpanded: case distance::DistanceType::InnerProduct: { float2 pvals; for (uint32_t i = threadIdx.x; i < dim; i += blockDim.x) { @@ -408,6 +409,7 @@ RAFT_KERNEL compute_similarity_kernel(uint32_t dim, diff -= pq_c; score += diff * diff; } break; + case distance::DistanceType::CosineExpanded: case distance::DistanceType::InnerProduct: { // NB: we negate the scores as we hardcoded select-topk to always compute the minimum float q; diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 5f812dc4f..65e8b6f44 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -38,6 +38,8 @@ #include #include #include +#include +#include #include #include #include @@ -110,6 +112,7 @@ void select_clusters(raft::resources const& handle, switch (metric) { case cuvs::distance::DistanceType::L2SqrtExpanded: case cuvs::distance::DistanceType::L2Expanded: norm_factor = 1.0 / -2.0; break; + case cuvs::distance::DistanceType::CosineExpanded: case cuvs::distance::DistanceType::InnerProduct: norm_factor = 0.0; break; default: RAFT_FAIL("Unsupported distance type %d.", int(metric)); } @@ -121,6 +124,10 @@ void select_clusters(raft::resources const& handle, uint32_t row = ix / dim_ext; return col < dim ? utils::mapping{}(queries[col + dim * row]) : norm_factor; }); + + auto float_queries_matrix_view = raft::make_device_matrix_view(float_queries, n_queries, dim_ext); + + raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_queries_matrix_view), float_queries_matrix_view, raft::linalg::NormType::L2Norm); float alpha; float beta; @@ -133,6 +140,7 @@ void select_clusters(raft::resources const& handle, gemm_k = dim + 1; RAFT_EXPECTS(gemm_k <= dim_ext, "unexpected gemm_k or dim_ext"); } break; + case cuvs::distance::DistanceType::CosineExpanded: case cuvs::distance::DistanceType::InnerProduct: { alpha = -1.0; beta = 0.0; @@ -363,6 +371,7 @@ void ivfpq_search_worker(raft::resources const& handle, // stores basediff (query[i] - center[i]) precomp_data_count = index.rot_dim(); } break; + case distance::DistanceType::CosineExpanded: case distance::DistanceType::InnerProduct: { // stores two components (query[i] * center[i], query[i] * center[i]) precomp_data_count = index.rot_dim() * 2; @@ -666,7 +675,7 @@ inline void search(raft::resources const& handle, uint32_t queries_batch = min(max_queries, n_queries - offset_q); raft::common::nvtx::range batch_scope( "ivf_pq::search-batch(queries: %u - %u)", offset_q, offset_q + queries_batch); - + select_clusters(handle, clusters_to_probe.data(), float_queries.data(), diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index e6d8efc93..dd5de99c9 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -888,6 +888,26 @@ inline auto enum_variety_l2sqrt() -> test_cases_t }); } +inline auto enum_variety_cosine() -> test_cases_t +{ + return map(enum_variety(), [](const ivf_pq_inputs& x) { + ivf_pq_inputs y(x); + if (y.min_recall.has_value()) { + if (y.search_params.lut_dtype == CUDA_R_8U) { + // InnerProduct score is signed, + // thus we're forced to used signed 8-bit representation, + // thus we have one bit less precision + y.min_recall = y.min_recall.value() * 0.90; + } else { + // In other cases it seems to perform a little bit better, still worse than L2 + y.min_recall = y.min_recall.value() * 0.94; + } + } + y.index_params.metric = distance::DistanceType::CosineExpanded; + return y; + }); +} + /** * Try different number of n_probes, some of which may trigger the non-fused version of the search * kernel. diff --git a/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu index cdc6c1b7e..47f914023 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu @@ -25,9 +25,9 @@ TEST_BUILD_HOST_INPUT_SEARCH(f32_f32_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_f32_i64) TEST_BUILD_EXTEND_SEARCH(f32_f32_i64) TEST_BUILD_SERIALIZE_SEARCH(f32_f32_i64) -INSTANTIATE(f32_f32_i64, defaults() + small_dims() + big_dims_moderate_lut()); +INSTANTIATE(f32_f32_i64, enum_variety_ip()); TEST_BUILD_SEARCH(f32_f32_i64_filter) -INSTANTIATE(f32_f32_i64_filter, defaults() + small_dims() + big_dims_moderate_lut()); +INSTANTIATE(f32_f32_i64_filter, enum_variety_ip()); } // namespace cuvs::neighbors::ivf_pq From 5eae823f61c30f686e6fc8066b7de203697cd050 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Mon, 23 Sep 2024 09:31:03 -0700 Subject: [PATCH 02/24] trial --- cpp/test/neighbors/ann_ivf_pq.cuh | 10 +++++----- cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu | 4 ++-- cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu | 4 ++-- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index dd5de99c9..e9539a999 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -128,8 +128,8 @@ void compare_vectors_l2( double d = dist(i); // The theoretical estimate of the error is hard to come up with, // the estimate below is based on experimentation + curse of dimensionality - ASSERT_LE(d, 1.2 * eps * std::pow(2.0, compression_ratio)) - << " (label = " << label << ", ix = " << i << ", eps = " << eps << ")"; + // ASSERT_LE(d, 1.2 * eps * std::pow(2.0, compression_ratio)) + // << " (label = " << label << ", ix = " << i << ", eps = " << eps << ")"; } } @@ -376,7 +376,7 @@ class ivf_pq_test : public ::testing::TestWithParam { // Pack a few vectors back to the list. int row_offset = 9; int n_vec = 3; - ASSERT_TRUE(row_offset + n_vec < n_rows); + // ASSERT_TRUE(row_offset + n_vec < n_rows); size_t offset = row_offset * index->pq_dim(); auto codes_to_pack = raft::make_device_matrix_view( codes.data_handle() + offset, n_vec, index->pq_dim()); @@ -390,7 +390,7 @@ class ivf_pq_test : public ::testing::TestWithParam { // Another test with the API that take list_data directly [[maybe_unused]] auto list_data = index->lists()[label]->data.view(); uint32_t n_take = 4; - ASSERT_TRUE(row_offset + n_take < n_rows); + // ASSERT_TRUE(row_offset + n_take < n_rows); auto codes2 = raft::make_device_matrix(handle_, n_take, index->pq_dim()); ivf_pq::helpers::codepacker::unpack( handle_, list_data, index->pq_bits(), row_offset, codes2.view()); @@ -874,7 +874,7 @@ inline auto enum_variety_ip() -> test_cases_t y.min_recall = y.min_recall.value() * 0.94; } } - y.index_params.metric = distance::DistanceType::InnerProduct; + y.index_params.metric = distance::DistanceType::CosineExpanded; return y; }); } diff --git a/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu index 80b0e2ccb..4f420177e 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu @@ -25,8 +25,8 @@ TEST_BUILD_SEARCH(f32_i08_i64) TEST_BUILD_HOST_INPUT_SEARCH(f32_i08_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_i08_i64) TEST_BUILD_SERIALIZE_SEARCH(f32_i08_i64) -INSTANTIATE(f32_i08_i64, defaults() + big_dims() + var_k()); +INSTANTIATE(f32_i08_i64, enum_variety_ip()); TEST_BUILD_SEARCH(f32_i08_i64_filter) -INSTANTIATE(f32_i08_i64_filter, defaults() + big_dims() + var_k()); +INSTANTIATE(f32_i08_i64_filter, enum_variety_ip()); } // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu index 0216a1e80..5e43dd781 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu @@ -25,8 +25,8 @@ TEST_BUILD_SEARCH(f32_u08_i64) TEST_BUILD_HOST_INPUT_SEARCH(f32_u08_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_u08_i64) TEST_BUILD_EXTEND_SEARCH(f32_u08_i64) -INSTANTIATE(f32_u08_i64, small_dims_per_cluster() + enum_variety()); +INSTANTIATE(f32_u08_i64, small_dims_per_cluster() + enum_variety_ip()); TEST_BUILD_SEARCH(f32_u08_i64_filter) -INSTANTIATE(f32_u08_i64_filter, small_dims_per_cluster() + enum_variety()); +INSTANTIATE(f32_u08_i64_filter, small_dims_per_cluster() + enum_variety_ip()); } // namespace cuvs::neighbors::ivf_pq From 0a860d4f0f33bda3a24c3a29397191a4b1846630 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Tue, 24 Sep 2024 15:26:13 -0700 Subject: [PATCH 03/24] debug --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 58 ++++++++++++++----- .../neighbors/ivf_pq/ivf_pq_build_common.cu | 2 +- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 53 ++++++++++++++--- cpp/test/CMakeLists.txt | 4 +- cpp/test/neighbors/ann_ivf_pq.cuh | 4 ++ .../ann_ivf_pq/test_float_int64_t.cu | 4 +- .../ann_ivf_pq/test_int8_t_int64_t.cu | 4 +- .../ann_ivf_pq/test_uint8_t_int64_t.cu | 4 +- cpp/test/neighbors/ann_utils.cuh | 2 + 9 files changed, 103 insertions(+), 32 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index 09e1396b3..dd62d2cca 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -256,6 +256,7 @@ void set_centers(raft::resources const& handle, index* index, const float* raft::linalg::L2Norm, true, stream); + RAFT_CUDA_TRY(cudaMemcpy2DAsync(index->centers().data_handle() + index->dim(), sizeof(float) * index->dim_ext(), center_norms.data(), @@ -1571,25 +1572,34 @@ void extend(raft::resources const& handle, stream)); vec_batches.prefetch_next_batch(); for (const auto& batch : vec_batches) { - // auto batch_data_view = raft::make_device_matrix_view( - // batch.data(), batch.size(), index->dim()); + auto batch_data_view = raft::make_device_matrix_view( + batch.data(), batch.size(), index->dim()); auto batch_labels_view = raft::make_device_vector_view( new_data_labels.data() + batch.offset(), batch.size()); - auto float_vec_batch = raft::make_device_matrix(handle, batch.size(), index->dim()); - raft::linalg::map_offset(handle, raft::make_device_vector_view(batch.data(), batch.size() * index->dim()), raft::make_device_vector_view(float_vec_batch.data_handle(), float_vec_batch.size()), [=]__device__(internal_extents_t idx, T i) {return utils::mapping{}(i);}); - if(index->metric() == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_vec_batch.view()), float_vec_batch.view(), raft::linalg::NormType::L2Norm); - } auto centers_view = raft::make_device_matrix_view( cluster_centers.data(), n_clusters, index->dim()); cuvs::cluster::kmeans::balanced_params kmeans_params; - kmeans_params.metric = cuvs::distance::DistanceType::InnerProduct; + if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { + auto float_vec_batch = raft::make_device_matrix(handle, batch.size(), index->dim()); + raft::linalg::map_offset(handle, raft::make_device_vector_view(batch.data(), batch.size() * index->dim()), raft::make_device_vector_view(float_vec_batch.data_handle(), float_vec_batch.size()), [=]__device__(internal_extents_t idx, T i) {return utils::mapping{}(i);}); + raft::print_device_vector("non_normalized_extend", batch.data(), index->dim(), std::cout); + raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_vec_batch.view()), float_vec_batch.view(), raft::linalg::NormType::L2Norm); + raft::print_device_vector("normalized_extend", float_vec_batch.data_handle(), index->dim(), std::cout); + kmeans_params.metric = distance::DistanceType::InnerProduct; cuvs::cluster::kmeans_balanced::predict(handle, kmeans_params, raft::make_const_mdspan(float_vec_batch.view()), centers_view, + batch_labels_view); + } else { + kmeans_params.metric = index->metric(); + cuvs::cluster::kmeans_balanced::predict(handle, + kmeans_params, + batch_data_view, + centers_view, batch_labels_view, utils::mapping{}); + } vec_batches.prefetch_next_batch(); // User needs to make sure kernel finishes its work before we overwrite batch in the next // iteration if different streams are used for kernel and copy. @@ -1639,11 +1649,21 @@ void extend(raft::resources const& handle, vec_batches.prefetch_next_batch(); for (const auto& vec_batch : vec_batches) { const auto& idx_batch = *idx_batches++; - auto float_vec_batch = raft::make_device_matrix(handle, vec_batch.size(), index->dim()); - raft::linalg::map_offset(handle, raft::make_device_vector_view(vec_batch.data(), vec_batch.size() * index->dim()), raft::make_device_vector_view(float_vec_batch.data_handle(), float_vec_batch.size()), [=]__device__(internal_extents_t idx, T i) {return utils::mapping{}(i);}); - if(index->metric() == cuvs::distance::DistanceType::CosineExpanded) { - raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_vec_batch.view()), float_vec_batch.view(), raft::linalg::NormType::L2Norm); - } + auto float_vec_batch = + raft::make_device_matrix(handle, vec_batch.size(), index->dim()); + raft::linalg::map_offset( + handle, + raft::make_device_vector_view(vec_batch.data(), + vec_batch.size() * index->dim()), + raft::make_device_vector_view(float_vec_batch.data_handle(), + vec_batch.size() * index->dim()), + [=] __device__(internal_extents_t idx, T i) { return utils::mapping{}(i); }); + if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { + raft::linalg::row_normalize(handle, + raft::make_const_mdspan(float_vec_batch.view()), + float_vec_batch.view(), + raft::linalg::NormType::L2Norm); + } process_and_fill_codes(handle, *index, float_vec_batch.data_handle(), @@ -1762,14 +1782,22 @@ auto build(raft::resources const& handle, // Train balanced hierarchical kmeans clustering auto trainset_const_view = raft::make_const_mdspan(trainset.view()); - raft::linalg::row_normalize(handle, trainset_const_view, trainset.view(), raft::linalg::NormType::L2Norm); auto centers_view = raft::make_device_matrix_view( cluster_centers, index.n_lists(), index.dim()); cuvs::cluster::kmeans::balanced_params kmeans_params; kmeans_params.n_iters = params.kmeans_n_iters; - kmeans_params.metric = cuvs::distance::DistanceType::InnerProduct; + if (index.metric() == distance::DistanceType::CosineExpanded) { + raft::print_device_vector("non_normalized_build", trainset.data_handle(), index.dim(), std::cout); + raft::linalg::row_normalize(handle, trainset_const_view, trainset.view(), raft::linalg::NormType::L2Norm); + raft::print_device_vector("normalized_build", trainset.data_handle(), index.dim(), std::cout); + kmeans_params.metric = distance::DistanceType::InnerProduct; + } else { + kmeans_params.metric = index.metric(); + } cuvs::cluster::kmeans_balanced::fit( handle, kmeans_params, trainset_const_view, centers_view, utils::mapping{}); + + // raft::linalg::row_normalize(handle, raft::make_const_mdspan(centers_view), centers_view, raft::linalg::NormType::L2Norm); // Trainset labels are needed for training PQ codebooks rmm::device_uvector labels(n_rows_train, stream, big_memory_resource); diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu b/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu index b34611c1d..b7dbafc28 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu @@ -242,7 +242,7 @@ void make_rotation_matrix(raft::resources const& handle, auto stream = raft::resource::get_cuda_stream(handle); bool inplace = n_rows == n_cols; uint32_t n = std::max(n_rows, n_cols); - if (force_random_rotation || !inplace) { + if (false) { rmm::device_uvector buf(inplace ? 0 : n * n, stream); float* mat = inplace ? rotation_matrix : buf.data(); raft::random::normal(handle, rng, mat, n * n, 0.0f, 1.0f); diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 65e8b6f44..40ed012ce 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -37,9 +37,9 @@ #include #include #include -#include #include #include +#include #include #include #include @@ -124,10 +124,6 @@ void select_clusters(raft::resources const& handle, uint32_t row = ix / dim_ext; return col < dim ? utils::mapping{}(queries[col + dim * row]) : norm_factor; }); - - auto float_queries_matrix_view = raft::make_device_matrix_view(float_queries, n_queries, dim_ext); - - raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_queries_matrix_view), float_queries_matrix_view, raft::linalg::NormType::L2Norm); float alpha; float beta; @@ -140,11 +136,22 @@ void select_clusters(raft::resources const& handle, gemm_k = dim + 1; RAFT_EXPECTS(gemm_k <= dim_ext, "unexpected gemm_k or dim_ext"); } break; - case cuvs::distance::DistanceType::CosineExpanded: case cuvs::distance::DistanceType::InnerProduct: { alpha = -1.0; beta = 0.0; } break; + case cuvs::distance::DistanceType::CosineExpanded: { + alpha = -1.0; + beta = 0.0; + + auto float_queries_matrix_view = + raft::make_device_matrix_view(float_queries, n_queries, dim_ext); + + raft::linalg::row_normalize(handle, + raft::make_const_mdspan(float_queries_matrix_view), + float_queries_matrix_view, + raft::linalg::NormType::L2Norm); + } break; default: RAFT_FAIL("Unsupported distance type %d.", int(metric)); } rmm::device_uvector qc_distances(n_queries * n_lists, stream, mr); @@ -164,6 +171,35 @@ void select_clusters(raft::resources const& handle, n_lists, stream); + if (metric == distance::DistanceType::CosineExpanded) { + // TODO: store dataset norms in a different manner for the cosine metric to avoid the copy here + auto center_norms = + raft::make_device_mdarray(handle, mr, raft::make_extents(n_lists)); + + cudaMemcpy2DAsync(center_norms.data_handle(), + sizeof(float), + cluster_centers + dim, + sizeof(float) * dim_ext, + sizeof(float), + n_lists, + cudaMemcpyDefault, + stream); + raft::linalg::map_offset( + handle, + raft::make_device_vector_view(center_norms.data_handle(), n_lists), + raft::sqrt_op{}); + + raft::linalg::matrixVectorOp(qc_distances.data(), + qc_distances.data(), + center_norms.data_handle(), + n_lists, + n_queries, + true, + true, + raft::div_checkzero_op{}, + stream); + } + // Select neighbor clusters for each query. rmm::device_uvector cluster_dists(n_queries * n_probes, stream, mr); cuvs::selection::select_k( @@ -373,7 +409,7 @@ void ivfpq_search_worker(raft::resources const& handle, } break; case distance::DistanceType::CosineExpanded: case distance::DistanceType::InnerProduct: { - // stores two components (query[i] * center[i], query[i] * center[i]) + // stores two components (query[i], query[i] * center[i]) precomp_data_count = index.rot_dim() * 2; } break; default: { @@ -517,6 +553,7 @@ struct ivfpq_search { { bool signed_metric = false; switch (metric) { + case cuvs::distance::DistanceType::CosineExpanded: signed_metric = true; break; case cuvs::distance::DistanceType::InnerProduct: signed_metric = true; break; default: break; } @@ -675,7 +712,7 @@ inline void search(raft::resources const& handle, uint32_t queries_batch = min(max_queries, n_queries - offset_q); raft::common::nvtx::range batch_scope( "ivf_pq::search-batch(queries: %u - %u)", offset_q, offset_q + queries_batch); - + select_clusters(handle, clusters_to_probe.data(), float_queries.data(), diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt index e04c39318..7c7166ffd 100644 --- a/cpp/test/CMakeLists.txt +++ b/cpp/test/CMakeLists.txt @@ -121,8 +121,8 @@ if(BUILD_TESTS) NEIGHBORS_ANN_IVF_PQ_TEST PATH neighbors/ann_ivf_pq/test_float_int64_t.cu - neighbors/ann_ivf_pq/test_int8_t_int64_t.cu - neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu + # neighbors/ann_ivf_pq/test_int8_t_int64_t.cu + # neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu GPUS 1 PERCENT diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index e9539a999..80f53df7b 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -24,6 +24,8 @@ #include #include #include +#include +#include #include #include #include @@ -168,6 +170,8 @@ class ivf_pq_test : public ::testing::TestWithParam { handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(0.1), DataT(2.0)); raft::random::uniform( handle_, r, search_queries.data(), ps.num_queries * ps.dim, DataT(0.1), DataT(2.0)); + // auto dv = raft::make_device_matrix_view(database.data(), (size_t)(ps.num_db_vecs), (size_t)ps.dim); + // raft::linalg::row_normalize(handle_, raft::make_const_mdspan(dv), dv, raft::linalg::NormType::L2Norm); } else { raft::random::uniformInt( handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(1), DataT(20)); diff --git a/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu index 47f914023..64c338906 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu @@ -25,9 +25,9 @@ TEST_BUILD_HOST_INPUT_SEARCH(f32_f32_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_f32_i64) TEST_BUILD_EXTEND_SEARCH(f32_f32_i64) TEST_BUILD_SERIALIZE_SEARCH(f32_f32_i64) -INSTANTIATE(f32_f32_i64, enum_variety_ip()); +INSTANTIATE(f32_f32_i64, enum_variety_cosine()); TEST_BUILD_SEARCH(f32_f32_i64_filter) -INSTANTIATE(f32_f32_i64_filter, enum_variety_ip()); +INSTANTIATE(f32_f32_i64_filter, enum_variety_cosine()); } // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu index 4f420177e..5c4154bc0 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu @@ -25,8 +25,8 @@ TEST_BUILD_SEARCH(f32_i08_i64) TEST_BUILD_HOST_INPUT_SEARCH(f32_i08_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_i08_i64) TEST_BUILD_SERIALIZE_SEARCH(f32_i08_i64) -INSTANTIATE(f32_i08_i64, enum_variety_ip()); +INSTANTIATE(f32_i08_i64, enum_variety_cosine()); TEST_BUILD_SEARCH(f32_i08_i64_filter) -INSTANTIATE(f32_i08_i64_filter, enum_variety_ip()); +INSTANTIATE(f32_i08_i64_filter, enum_variety_cosine()); } // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu index 5e43dd781..3b38fe42b 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu @@ -25,8 +25,8 @@ TEST_BUILD_SEARCH(f32_u08_i64) TEST_BUILD_HOST_INPUT_SEARCH(f32_u08_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_u08_i64) TEST_BUILD_EXTEND_SEARCH(f32_u08_i64) -INSTANTIATE(f32_u08_i64, small_dims_per_cluster() + enum_variety_ip()); +INSTANTIATE(f32_u08_i64, small_dims_per_cluster() + enum_variety_cosine()); TEST_BUILD_SEARCH(f32_u08_i64_filter) -INSTANTIATE(f32_u08_i64_filter, small_dims_per_cluster() + enum_variety_ip()); +INSTANTIATE(f32_u08_i64_filter, small_dims_per_cluster() + enum_variety_cosine()); } // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/test/neighbors/ann_utils.cuh b/cpp/test/neighbors/ann_utils.cuh index b08e1d725..fd81a8c97 100644 --- a/cpp/test/neighbors/ann_utils.cuh +++ b/cpp/test/neighbors/ann_utils.cuh @@ -269,6 +269,8 @@ auto eval_neighbours(const std::vector& expected_idx, auto [actual_recall, match_count, total_count] = calc_recall(expected_idx, actual_idx, expected_dist, actual_dist, rows, cols, eps); double error_margin = (actual_recall - min_recall) / std::max(1.0 - min_recall, eps); + raft::print_host_vector("expected_dist", expected_dist.data(), 100, std::cout); + raft::print_host_vector("actual_dist", actual_dist.data(), 100, std::cout); RAFT_LOG_INFO("Recall = %f (%zu/%zu), the error is %2.1f%% %s the threshold (eps = %f).", actual_recall, From 55c17fd4dc4fe8644536b8d33d9cc77b73e9bd9d Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Tue, 24 Sep 2024 16:57:14 -0700 Subject: [PATCH 04/24] debug --- cpp/src/neighbors/ivf_common.cuh | 9 ++++++++- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 2 +- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 7 ++++++- 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/cpp/src/neighbors/ivf_common.cuh b/cpp/src/neighbors/ivf_common.cuh index fb73fb8a9..0027d494f 100644 --- a/cpp/src/neighbors/ivf_common.cuh +++ b/cpp/src/neighbors/ivf_common.cuh @@ -254,7 +254,14 @@ void postprocess_distances(ScoreOutT* out, // [n_queries, topk] raft::linalg::unaryOp(out, in, len, raft::sqrt_op{}, stream); } } break; - case distance::DistanceType::CosineExpanded: + case distance::DistanceType::CosineExpanded: { + raft::linalg::unaryOp( + out, + in, + len, + raft::compose_op(raft::add_const_op{1.0}, raft::cast_op{}), + stream); + } break; case distance::DistanceType::InnerProduct: { float factor = (account_for_max_close ? -1.0 : 1.0) * scaling_factor * scaling_factor; if (factor != 1.0) { diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index dd62d2cca..3bf48b079 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -1592,7 +1592,7 @@ void extend(raft::resources const& handle, centers_view, batch_labels_view); } else { - kmeans_params.metric = index->metric(); + kmeans_params.metric = static_cast((int)index->metric()); cuvs::cluster::kmeans_balanced::predict(handle, kmeans_params, batch_data_view, diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 40ed012ce..393ae29f1 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -38,6 +38,7 @@ #include #include #include +#include #include #include #include @@ -744,7 +745,11 @@ inline void search(raft::resources const& handle, rot_queries.data(), index.rot_dim(), stream); - + + raft::print_device_vector("rot_queries", rot_queries.data(), index.rot_dim(), std::cout); + auto rot_queries_view = raft::make_device_matrix_view(rot_queries.data(), max_queries, index.rot_dim()); + raft::linalg::row_normalize(handle, raft::make_const_mdspan(rot_queries_view), rot_queries_view, raft::linalg::NormType::L2Norm); + raft::print_device_vector("rot_queries_normalized", rot_queries.data(), index.rot_dim(), std::cout); for (uint32_t offset_b = 0; offset_b < queries_batch; offset_b += max_batch_size) { uint32_t batch_size = min(max_batch_size, queries_batch - offset_b); /* The distance calculation is done in the rotated/transformed space; From e3490e3582bc089ec0d287b47294902355fc4a25 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Tue, 24 Sep 2024 17:36:47 -0700 Subject: [PATCH 05/24] undo change --- cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu | 2 +- cpp/test/neighbors/ann_ivf_pq.cuh | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu b/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu index b7dbafc28..b34611c1d 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build_common.cu @@ -242,7 +242,7 @@ void make_rotation_matrix(raft::resources const& handle, auto stream = raft::resource::get_cuda_stream(handle); bool inplace = n_rows == n_cols; uint32_t n = std::max(n_rows, n_cols); - if (false) { + if (force_random_rotation || !inplace) { rmm::device_uvector buf(inplace ? 0 : n * n, stream); float* mat = inplace ? rotation_matrix : buf.data(); raft::random::normal(handle, rng, mat, n * n, 0.0f, 1.0f); diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index 80f53df7b..6c711ede7 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -898,6 +899,7 @@ inline auto enum_variety_cosine() -> test_cases_t ivf_pq_inputs y(x); if (y.min_recall.has_value()) { if (y.search_params.lut_dtype == CUDA_R_8U) { + y.search_params.lut_dtype = CUDA_R_16F; // InnerProduct score is signed, // thus we're forced to used signed 8-bit representation, // thus we have one bit less precision From be343be531aa9a732c2f4e26de6702326e4b5f88 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Tue, 24 Sep 2024 17:44:18 -0700 Subject: [PATCH 06/24] style --- cpp/src/neighbors/ivf_common.cuh | 10 ++-- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 58 ++++++++++++++-------- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 13 +++-- cpp/test/neighbors/ann_ivf_pq.cuh | 5 +- 4 files changed, 53 insertions(+), 33 deletions(-) diff --git a/cpp/src/neighbors/ivf_common.cuh b/cpp/src/neighbors/ivf_common.cuh index 0027d494f..05035e14d 100644 --- a/cpp/src/neighbors/ivf_common.cuh +++ b/cpp/src/neighbors/ivf_common.cuh @@ -256,11 +256,11 @@ void postprocess_distances(ScoreOutT* out, // [n_queries, topk] } break; case distance::DistanceType::CosineExpanded: { raft::linalg::unaryOp( - out, - in, - len, - raft::compose_op(raft::add_const_op{1.0}, raft::cast_op{}), - stream); + out, + in, + len, + raft::compose_op(raft::add_const_op{1.0}, raft::cast_op{}), + stream); } break; case distance::DistanceType::InnerProduct: { float factor = (account_for_max_close ? -1.0 : 1.0) * scaling_factor * scaling_factor; diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index 3bf48b079..b403b65b8 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -1580,25 +1580,36 @@ void extend(raft::resources const& handle, cluster_centers.data(), n_clusters, index->dim()); cuvs::cluster::kmeans::balanced_params kmeans_params; if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { - auto float_vec_batch = raft::make_device_matrix(handle, batch.size(), index->dim()); - raft::linalg::map_offset(handle, raft::make_device_vector_view(batch.data(), batch.size() * index->dim()), raft::make_device_vector_view(float_vec_batch.data_handle(), float_vec_batch.size()), [=]__device__(internal_extents_t idx, T i) {return utils::mapping{}(i);}); - raft::print_device_vector("non_normalized_extend", batch.data(), index->dim(), std::cout); - raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_vec_batch.view()), float_vec_batch.view(), raft::linalg::NormType::L2Norm); - raft::print_device_vector("normalized_extend", float_vec_batch.data_handle(), index->dim(), std::cout); + auto float_vec_batch = + raft::make_device_matrix(handle, batch.size(), index->dim()); + raft::linalg::map_offset( + handle, + raft::make_device_vector_view(batch.data(), + batch.size() * index->dim()), + raft::make_device_vector_view(float_vec_batch.data_handle(), + float_vec_batch.size()), + [=] __device__(internal_extents_t idx, T i) { return utils::mapping{}(i); }); + raft::print_device_vector("non_normalized_extend", batch.data(), index->dim(), std::cout); + raft::linalg::row_normalize(handle, + raft::make_const_mdspan(float_vec_batch.view()), + float_vec_batch.view(), + raft::linalg::NormType::L2Norm); + raft::print_device_vector( + "normalized_extend", float_vec_batch.data_handle(), index->dim(), std::cout); kmeans_params.metric = distance::DistanceType::InnerProduct; - cuvs::cluster::kmeans_balanced::predict(handle, - kmeans_params, - raft::make_const_mdspan(float_vec_batch.view()), - centers_view, - batch_labels_view); + cuvs::cluster::kmeans_balanced::predict(handle, + kmeans_params, + raft::make_const_mdspan(float_vec_batch.view()), + centers_view, + batch_labels_view); } else { - kmeans_params.metric = static_cast((int)index->metric()); - cuvs::cluster::kmeans_balanced::predict(handle, - kmeans_params, - batch_data_view, - centers_view, - batch_labels_view, - utils::mapping{}); + kmeans_params.metric = static_cast((int)index->metric()); + cuvs::cluster::kmeans_balanced::predict(handle, + kmeans_params, + batch_data_view, + centers_view, + batch_labels_view, + utils::mapping{}); } vec_batches.prefetch_next_batch(); // User needs to make sure kernel finishes its work before we overwrite batch in the next @@ -1787,17 +1798,20 @@ auto build(raft::resources const& handle, cuvs::cluster::kmeans::balanced_params kmeans_params; kmeans_params.n_iters = params.kmeans_n_iters; if (index.metric() == distance::DistanceType::CosineExpanded) { - raft::print_device_vector("non_normalized_build", trainset.data_handle(), index.dim(), std::cout); - raft::linalg::row_normalize(handle, trainset_const_view, trainset.view(), raft::linalg::NormType::L2Norm); + raft::print_device_vector( + "non_normalized_build", trainset.data_handle(), index.dim(), std::cout); + raft::linalg::row_normalize( + handle, trainset_const_view, trainset.view(), raft::linalg::NormType::L2Norm); raft::print_device_vector("normalized_build", trainset.data_handle(), index.dim(), std::cout); - kmeans_params.metric = distance::DistanceType::InnerProduct; + kmeans_params.metric = distance::DistanceType::InnerProduct; } else { kmeans_params.metric = index.metric(); } cuvs::cluster::kmeans_balanced::fit( handle, kmeans_params, trainset_const_view, centers_view, utils::mapping{}); - - // raft::linalg::row_normalize(handle, raft::make_const_mdspan(centers_view), centers_view, raft::linalg::NormType::L2Norm); + + // raft::linalg::row_normalize(handle, raft::make_const_mdspan(centers_view), centers_view, + // raft::linalg::NormType::L2Norm); // Trainset labels are needed for training PQ codebooks rmm::device_uvector labels(n_rows_train, stream, big_memory_resource); diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 393ae29f1..4b717db70 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -745,11 +745,16 @@ inline void search(raft::resources const& handle, rot_queries.data(), index.rot_dim(), stream); - + raft::print_device_vector("rot_queries", rot_queries.data(), index.rot_dim(), std::cout); - auto rot_queries_view = raft::make_device_matrix_view(rot_queries.data(), max_queries, index.rot_dim()); - raft::linalg::row_normalize(handle, raft::make_const_mdspan(rot_queries_view), rot_queries_view, raft::linalg::NormType::L2Norm); - raft::print_device_vector("rot_queries_normalized", rot_queries.data(), index.rot_dim(), std::cout); + auto rot_queries_view = raft::make_device_matrix_view( + rot_queries.data(), max_queries, index.rot_dim()); + raft::linalg::row_normalize(handle, + raft::make_const_mdspan(rot_queries_view), + rot_queries_view, + raft::linalg::NormType::L2Norm); + raft::print_device_vector( + "rot_queries_normalized", rot_queries.data(), index.rot_dim(), std::cout); for (uint32_t offset_b = 0; offset_b < queries_batch; offset_b += max_batch_size) { uint32_t batch_size = min(max_batch_size, queries_batch - offset_b); /* The distance calculation is done in the rotated/transformed space; diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index 6c711ede7..e458f5eb7 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -171,8 +171,9 @@ class ivf_pq_test : public ::testing::TestWithParam { handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(0.1), DataT(2.0)); raft::random::uniform( handle_, r, search_queries.data(), ps.num_queries * ps.dim, DataT(0.1), DataT(2.0)); - // auto dv = raft::make_device_matrix_view(database.data(), (size_t)(ps.num_db_vecs), (size_t)ps.dim); - // raft::linalg::row_normalize(handle_, raft::make_const_mdspan(dv), dv, raft::linalg::NormType::L2Norm); + // auto dv = raft::make_device_matrix_view(database.data(), + // (size_t)(ps.num_db_vecs), (size_t)ps.dim); raft::linalg::row_normalize(handle_, + // raft::make_const_mdspan(dv), dv, raft::linalg::NormType::L2Norm); } else { raft::random::uniformInt( handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(1), DataT(20)); From f3b50e404f22b18cd9c1cd40969c0656af16dc6c Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 25 Sep 2024 12:52:27 -0700 Subject: [PATCH 07/24] tests passing: --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 10 +++++---- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 2 +- cpp/test/neighbors/ann_ivf_pq.cuh | 17 ++++++++++---- python/cuvs/CMakeLists.txt | 22 +++++++++++++------ .../cuvs/neighbors/filters/CMakeLists.txt | 3 +-- 5 files changed, 36 insertions(+), 18 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index b403b65b8..570439d2a 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -1579,6 +1579,7 @@ void extend(raft::resources const& handle, auto centers_view = raft::make_device_matrix_view( cluster_centers.data(), n_clusters, index->dim()); cuvs::cluster::kmeans::balanced_params kmeans_params; + raft::print_device_vector("non_normalized_extend", batch.data(), index->dim(), std::cout); if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { auto float_vec_batch = raft::make_device_matrix(handle, batch.size(), index->dim()); @@ -1589,7 +1590,6 @@ void extend(raft::resources const& handle, raft::make_device_vector_view(float_vec_batch.data_handle(), float_vec_batch.size()), [=] __device__(internal_extents_t idx, T i) { return utils::mapping{}(i); }); - raft::print_device_vector("non_normalized_extend", batch.data(), index->dim(), std::cout); raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_vec_batch.view()), float_vec_batch.view(), @@ -1798,17 +1798,19 @@ auto build(raft::resources const& handle, cuvs::cluster::kmeans::balanced_params kmeans_params; kmeans_params.n_iters = params.kmeans_n_iters; if (index.metric() == distance::DistanceType::CosineExpanded) { - raft::print_device_vector( - "non_normalized_build", trainset.data_handle(), index.dim(), std::cout); + // raft::print_device_vector( + // "non_normalized_build", trainset.data_handle(), index.dim(), std::cout); raft::linalg::row_normalize( handle, trainset_const_view, trainset.view(), raft::linalg::NormType::L2Norm); raft::print_device_vector("normalized_build", trainset.data_handle(), index.dim(), std::cout); kmeans_params.metric = distance::DistanceType::InnerProduct; } else { - kmeans_params.metric = index.metric(); + kmeans_params.metric = static_cast((int)index.metric()); } cuvs::cluster::kmeans_balanced::fit( handle, kmeans_params, trainset_const_view, centers_view, utils::mapping{}); + + raft::print_device_vector("centers", centers_view.data_handle(), index.dim(), std::cout); // raft::linalg::row_normalize(handle, raft::make_const_mdspan(centers_view), centers_view, // raft::linalg::NormType::L2Norm); diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 4b717db70..f52c15731 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -172,7 +172,7 @@ void select_clusters(raft::resources const& handle, n_lists, stream); - if (metric == distance::DistanceType::CosineExpanded) { + if (false) { // TODO: store dataset norms in a different manner for the cosine metric to avoid the copy here auto center_norms = raft::make_device_mdarray(handle, mr, raft::make_extents(n_lists)); diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index e458f5eb7..5fbf5d802 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -172,8 +172,11 @@ class ivf_pq_test : public ::testing::TestWithParam { raft::random::uniform( handle_, r, search_queries.data(), ps.num_queries * ps.dim, DataT(0.1), DataT(2.0)); // auto dv = raft::make_device_matrix_view(database.data(), - // (size_t)(ps.num_db_vecs), (size_t)ps.dim); raft::linalg::row_normalize(handle_, - // raft::make_const_mdspan(dv), dv, raft::linalg::NormType::L2Norm); + // (size_t)(ps.num_db_vecs), (size_t)ps.dim); + // raft::linalg::row_normalize(handle_, raft::make_const_mdspan(dv), dv, raft::linalg::NormType::L2Norm); + // auto sv = raft::make_device_matrix_view(search_queries.data(), + // (size_t)(ps.num_db_vecs), (size_t)ps.dim); + // raft::linalg::row_normalize(handle_, raft::make_const_mdspan(sv), sv, raft::linalg::NormType::L2Norm); } else { raft::random::uniformInt( handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(1), DataT(20)); @@ -198,7 +201,7 @@ class ivf_pq_test : public ::testing::TestWithParam { ps.num_db_vecs, ps.dim, ps.k, - static_cast((int)ps.index_params.metric)); + cuvs::distance::DistanceType::CosineExpanded); distances_ref.resize(queries_size); raft::update_host(distances_ref.data(), distances_naive_dev.data(), queries_size, stream_); indices_ref.resize(queries_size); @@ -556,6 +559,12 @@ class ivf_pq_filter_test : public ::testing::TestWithParam { handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(0.1), DataT(2.0)); raft::random::uniform( handle_, r, search_queries.data(), ps.num_queries * ps.dim, DataT(0.1), DataT(2.0)); + // auto dv = raft::make_device_matrix_view(database.data(), + // (size_t)(ps.num_db_vecs), (size_t)ps.dim); + // raft::linalg::row_normalize(handle_, raft::make_const_mdspan(dv), dv, raft::linalg::NormType::L2Norm); + // auto sv = raft::make_device_matrix_view(search_queries.data(), + // (size_t)(ps.num_db_vecs), (size_t)ps.dim); + // raft::linalg::row_normalize(handle_, raft::make_const_mdspan(sv), sv, raft::linalg::NormType::L2Norm); } else { raft::random::uniformInt( handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(1), DataT(20)); @@ -580,7 +589,7 @@ class ivf_pq_filter_test : public ::testing::TestWithParam { ps.num_db_vecs - test_ivf_sample_filter::offset, ps.dim, ps.k, - static_cast((int)ps.index_params.metric)); + cuvs::distance::DistanceType::CosineExpanded); raft::linalg::addScalar(indices_naive_dev.data(), indices_naive_dev.data(), IdxT(test_ivf_sample_filter::offset), diff --git a/python/cuvs/CMakeLists.txt b/python/cuvs/CMakeLists.txt index 7d2f8dcf9..feb3bd58c 100644 --- a/python/cuvs/CMakeLists.txt +++ b/python/cuvs/CMakeLists.txt @@ -83,14 +83,22 @@ if(NOT cuvs_FOUND) if(NOT CUDA_STATIC_MATH_LIBRARIES AND USE_CUDA_MATH_WHEELS) set(rpaths - "$ORIGIN/../nvidia/cublas/lib" - "$ORIGIN/../nvidia/curand/lib" - "$ORIGIN/../nvidia/cusolver/lib" - "$ORIGIN/../nvidia/cusparse/lib" - "$ORIGIN/../nvidia/nvjitlink/lib" + "$ORIGIN/../nvidia/cublas/lib" + "$ORIGIN/../nvidia/curand/lib" + "$ORIGIN/../nvidia/cusolver/lib" + "$ORIGIN/../nvidia/cusparse/lib" + "$ORIGIN/../nvidia/nvjitlink/lib" + ) + set_property( + TARGET cuvs + PROPERTY INSTALL_RPATH ${rpaths} + APPEND + ) + set_property( + TARGET cuvs_c + PROPERTY INSTALL_RPATH ${rpaths} + APPEND ) - set_property(TARGET cuvs PROPERTY INSTALL_RPATH ${rpaths} APPEND) - set_property(TARGET cuvs_c PROPERTY INSTALL_RPATH ${rpaths} APPEND) endif() set(cython_lib_dir cuvs) diff --git a/python/cuvs/cuvs/neighbors/filters/CMakeLists.txt b/python/cuvs/cuvs/neighbors/filters/CMakeLists.txt index 8f281d1c8..c90615feb 100644 --- a/python/cuvs/cuvs/neighbors/filters/CMakeLists.txt +++ b/python/cuvs/cuvs/neighbors/filters/CMakeLists.txt @@ -20,6 +20,5 @@ set(linked_libraries cuvs::cuvs cuvs::c_api) rapids_cython_create_modules( CXX SOURCE_FILES "${cython_sources}" - LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS cuvs MODULE_PREFIX - neighbors_prefilter_ + LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS cuvs MODULE_PREFIX neighbors_prefilter_ ) From 3967c4c308c85e4ca50a35dbff515c877bbb6fdb Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 25 Sep 2024 12:55:33 -0700 Subject: [PATCH 08/24] remove debug statements --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 13 ++----------- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 4 +--- cpp/test/neighbors/ann_utils.cuh | 2 -- 3 files changed, 3 insertions(+), 16 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index 570439d2a..a5fc76ae0 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -1579,7 +1579,7 @@ void extend(raft::resources const& handle, auto centers_view = raft::make_device_matrix_view( cluster_centers.data(), n_clusters, index->dim()); cuvs::cluster::kmeans::balanced_params kmeans_params; - raft::print_device_vector("non_normalized_extend", batch.data(), index->dim(), std::cout); + if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { auto float_vec_batch = raft::make_device_matrix(handle, batch.size(), index->dim()); @@ -1594,8 +1594,7 @@ void extend(raft::resources const& handle, raft::make_const_mdspan(float_vec_batch.view()), float_vec_batch.view(), raft::linalg::NormType::L2Norm); - raft::print_device_vector( - "normalized_extend", float_vec_batch.data_handle(), index->dim(), std::cout); + kmeans_params.metric = distance::DistanceType::InnerProduct; cuvs::cluster::kmeans_balanced::predict(handle, kmeans_params, @@ -1798,11 +1797,8 @@ auto build(raft::resources const& handle, cuvs::cluster::kmeans::balanced_params kmeans_params; kmeans_params.n_iters = params.kmeans_n_iters; if (index.metric() == distance::DistanceType::CosineExpanded) { - // raft::print_device_vector( - // "non_normalized_build", trainset.data_handle(), index.dim(), std::cout); raft::linalg::row_normalize( handle, trainset_const_view, trainset.view(), raft::linalg::NormType::L2Norm); - raft::print_device_vector("normalized_build", trainset.data_handle(), index.dim(), std::cout); kmeans_params.metric = distance::DistanceType::InnerProduct; } else { kmeans_params.metric = static_cast((int)index.metric()); @@ -1810,11 +1806,6 @@ auto build(raft::resources const& handle, cuvs::cluster::kmeans_balanced::fit( handle, kmeans_params, trainset_const_view, centers_view, utils::mapping{}); - raft::print_device_vector("centers", centers_view.data_handle(), index.dim(), std::cout); - - // raft::linalg::row_normalize(handle, raft::make_const_mdspan(centers_view), centers_view, - // raft::linalg::NormType::L2Norm); - // Trainset labels are needed for training PQ codebooks rmm::device_uvector labels(n_rows_train, stream, big_memory_resource); auto centers_const_view = raft::make_device_matrix_view( diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index f52c15731..b4893f31d 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -746,15 +746,13 @@ inline void search(raft::resources const& handle, index.rot_dim(), stream); - raft::print_device_vector("rot_queries", rot_queries.data(), index.rot_dim(), std::cout); auto rot_queries_view = raft::make_device_matrix_view( rot_queries.data(), max_queries, index.rot_dim()); raft::linalg::row_normalize(handle, raft::make_const_mdspan(rot_queries_view), rot_queries_view, raft::linalg::NormType::L2Norm); - raft::print_device_vector( - "rot_queries_normalized", rot_queries.data(), index.rot_dim(), std::cout); + for (uint32_t offset_b = 0; offset_b < queries_batch; offset_b += max_batch_size) { uint32_t batch_size = min(max_batch_size, queries_batch - offset_b); /* The distance calculation is done in the rotated/transformed space; diff --git a/cpp/test/neighbors/ann_utils.cuh b/cpp/test/neighbors/ann_utils.cuh index fd81a8c97..b08e1d725 100644 --- a/cpp/test/neighbors/ann_utils.cuh +++ b/cpp/test/neighbors/ann_utils.cuh @@ -269,8 +269,6 @@ auto eval_neighbours(const std::vector& expected_idx, auto [actual_recall, match_count, total_count] = calc_recall(expected_idx, actual_idx, expected_dist, actual_dist, rows, cols, eps); double error_margin = (actual_recall - min_recall) / std::max(1.0 - min_recall, eps); - raft::print_host_vector("expected_dist", expected_dist.data(), 100, std::cout); - raft::print_host_vector("actual_dist", actual_dist.data(), 100, std::cout); RAFT_LOG_INFO("Recall = %f (%zu/%zu), the error is %2.1f%% %s the threshold (eps = %f).", actual_recall, From 442d65e2bc52f73abc2fda2f67168d3b246579f0 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 25 Sep 2024 14:41:53 -0700 Subject: [PATCH 09/24] add assertions --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 6 ++- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 30 ------------- cpp/test/CMakeLists.txt | 4 +- cpp/test/neighbors/ann_ivf_pq.cuh | 43 ++++++++----------- .../ann_ivf_pq/test_float_int64_t.cu | 4 +- .../ann_ivf_pq/test_int8_t_int64_t.cu | 4 +- .../ann_ivf_pq/test_uint8_t_int64_t.cu | 4 +- 7 files changed, 31 insertions(+), 64 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index a5fc76ae0..ae215bddc 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -30,6 +30,7 @@ #include "../../cluster/kmeans_balanced.cuh" #include +#include #include #include #include @@ -1725,6 +1726,9 @@ auto build(raft::resources const& handle, << (int)params.pq_dim << std::endl; RAFT_EXPECTS(n_rows > 0 && dim > 0, "empty dataset"); RAFT_EXPECTS(n_rows >= params.n_lists, "number of rows can't be less than n_lists"); + if (params.metric == cuvs::distance::DistanceType::CosineExpanded && params.codebook_kind == codebook_gen::PER_CLUSTER) { + RAFT_FAIL("CosineExpanded metric only supported for codebook_gen::PER_SUBSPACE"); + } auto stream = raft::resource::get_cuda_stream(handle); @@ -1805,7 +1809,7 @@ auto build(raft::resources const& handle, } cuvs::cluster::kmeans_balanced::fit( handle, kmeans_params, trainset_const_view, centers_view, utils::mapping{}); - + // Trainset labels are needed for training PQ codebooks rmm::device_uvector labels(n_rows_train, stream, big_memory_resource); auto centers_const_view = raft::make_device_matrix_view( diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index b4893f31d..b118d52ec 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -172,35 +172,6 @@ void select_clusters(raft::resources const& handle, n_lists, stream); - if (false) { - // TODO: store dataset norms in a different manner for the cosine metric to avoid the copy here - auto center_norms = - raft::make_device_mdarray(handle, mr, raft::make_extents(n_lists)); - - cudaMemcpy2DAsync(center_norms.data_handle(), - sizeof(float), - cluster_centers + dim, - sizeof(float) * dim_ext, - sizeof(float), - n_lists, - cudaMemcpyDefault, - stream); - raft::linalg::map_offset( - handle, - raft::make_device_vector_view(center_norms.data_handle(), n_lists), - raft::sqrt_op{}); - - raft::linalg::matrixVectorOp(qc_distances.data(), - qc_distances.data(), - center_norms.data_handle(), - n_lists, - n_queries, - true, - true, - raft::div_checkzero_op{}, - stream); - } - // Select neighbor clusters for each query. rmm::device_uvector cluster_dists(n_queries * n_probes, stream, mr); cuvs::selection::select_k( @@ -752,7 +723,6 @@ inline void search(raft::resources const& handle, raft::make_const_mdspan(rot_queries_view), rot_queries_view, raft::linalg::NormType::L2Norm); - for (uint32_t offset_b = 0; offset_b < queries_batch; offset_b += max_batch_size) { uint32_t batch_size = min(max_batch_size, queries_batch - offset_b); /* The distance calculation is done in the rotated/transformed space; diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt index 7c7166ffd..e04c39318 100644 --- a/cpp/test/CMakeLists.txt +++ b/cpp/test/CMakeLists.txt @@ -121,8 +121,8 @@ if(BUILD_TESTS) NEIGHBORS_ANN_IVF_PQ_TEST PATH neighbors/ann_ivf_pq/test_float_int64_t.cu - # neighbors/ann_ivf_pq/test_int8_t_int64_t.cu - # neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu + neighbors/ann_ivf_pq/test_int8_t_int64_t.cu + neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu GPUS 1 PERCENT diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index 5fbf5d802..0630e1a8a 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -21,12 +21,10 @@ #include #include -#include +#include #include #include #include -#include -#include #include #include #include @@ -131,8 +129,8 @@ void compare_vectors_l2( double d = dist(i); // The theoretical estimate of the error is hard to come up with, // the estimate below is based on experimentation + curse of dimensionality - // ASSERT_LE(d, 1.2 * eps * std::pow(2.0, compression_ratio)) - // << " (label = " << label << ", ix = " << i << ", eps = " << eps << ")"; + ASSERT_LE(d, 1.2 * eps * std::pow(2.0, compression_ratio)) + << " (label = " << label << ", ix = " << i << ", eps = " << eps << ")"; } } @@ -171,12 +169,6 @@ class ivf_pq_test : public ::testing::TestWithParam { handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(0.1), DataT(2.0)); raft::random::uniform( handle_, r, search_queries.data(), ps.num_queries * ps.dim, DataT(0.1), DataT(2.0)); - // auto dv = raft::make_device_matrix_view(database.data(), - // (size_t)(ps.num_db_vecs), (size_t)ps.dim); - // raft::linalg::row_normalize(handle_, raft::make_const_mdspan(dv), dv, raft::linalg::NormType::L2Norm); - // auto sv = raft::make_device_matrix_view(search_queries.data(), - // (size_t)(ps.num_db_vecs), (size_t)ps.dim); - // raft::linalg::row_normalize(handle_, raft::make_const_mdspan(sv), sv, raft::linalg::NormType::L2Norm); } else { raft::random::uniformInt( handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(1), DataT(20)); @@ -201,7 +193,7 @@ class ivf_pq_test : public ::testing::TestWithParam { ps.num_db_vecs, ps.dim, ps.k, - cuvs::distance::DistanceType::CosineExpanded); + static_cast((int)ps.index_params.metric)); distances_ref.resize(queries_size); raft::update_host(distances_ref.data(), distances_naive_dev.data(), queries_size, stream_); indices_ref.resize(queries_size); @@ -291,6 +283,10 @@ class ivf_pq_test : public ::testing::TestWithParam { uint32_t n_take, uint32_t n_skip) { + // the original data cannot be reconstructed since the dataset was normalized + if (index.metric() == cuvs::distance::DistanceType::CosineExpanded) { + return; + } auto& rec_list = index.lists()[label]; auto dim = index.dim(); n_take = std::min(n_take, rec_list->size.load()); @@ -322,6 +318,9 @@ class ivf_pq_test : public ::testing::TestWithParam { auto old_list = index->lists()[label]; auto n_rows = old_list->size.load(); if (n_rows == 0) { return; } + if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { + return; + } auto vectors_1 = raft::make_device_matrix(handle_, n_rows, index->dim()); auto indices = raft::make_device_vector(handle_, n_rows); @@ -385,7 +384,7 @@ class ivf_pq_test : public ::testing::TestWithParam { // Pack a few vectors back to the list. int row_offset = 9; int n_vec = 3; - // ASSERT_TRUE(row_offset + n_vec < n_rows); + ASSERT_TRUE(row_offset + n_vec < n_rows); size_t offset = row_offset * index->pq_dim(); auto codes_to_pack = raft::make_device_matrix_view( codes.data_handle() + offset, n_vec, index->pq_dim()); @@ -399,7 +398,7 @@ class ivf_pq_test : public ::testing::TestWithParam { // Another test with the API that take list_data directly [[maybe_unused]] auto list_data = index->lists()[label]->data.view(); uint32_t n_take = 4; - // ASSERT_TRUE(row_offset + n_take < n_rows); + ASSERT_TRUE(row_offset + n_take < n_rows); auto codes2 = raft::make_device_matrix(handle_, n_take, index->pq_dim()); ivf_pq::helpers::codepacker::unpack( handle_, list_data, index->pq_bits(), row_offset, codes2.view()); @@ -559,12 +558,6 @@ class ivf_pq_filter_test : public ::testing::TestWithParam { handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(0.1), DataT(2.0)); raft::random::uniform( handle_, r, search_queries.data(), ps.num_queries * ps.dim, DataT(0.1), DataT(2.0)); - // auto dv = raft::make_device_matrix_view(database.data(), - // (size_t)(ps.num_db_vecs), (size_t)ps.dim); - // raft::linalg::row_normalize(handle_, raft::make_const_mdspan(dv), dv, raft::linalg::NormType::L2Norm); - // auto sv = raft::make_device_matrix_view(search_queries.data(), - // (size_t)(ps.num_db_vecs), (size_t)ps.dim); - // raft::linalg::row_normalize(handle_, raft::make_const_mdspan(sv), sv, raft::linalg::NormType::L2Norm); } else { raft::random::uniformInt( handle_, r, database.data(), ps.num_db_vecs * ps.dim, DataT(1), DataT(20)); @@ -589,7 +582,7 @@ class ivf_pq_filter_test : public ::testing::TestWithParam { ps.num_db_vecs - test_ivf_sample_filter::offset, ps.dim, ps.k, - cuvs::distance::DistanceType::CosineExpanded); + static_cast((int)ps.index_params.metric)); raft::linalg::addScalar(indices_naive_dev.data(), indices_naive_dev.data(), IdxT(test_ivf_sample_filter::offset), @@ -889,7 +882,7 @@ inline auto enum_variety_ip() -> test_cases_t y.min_recall = y.min_recall.value() * 0.94; } } - y.index_params.metric = distance::DistanceType::CosineExpanded; + y.index_params.metric = distance::DistanceType::InnerProduct; return y; }); } @@ -909,17 +902,17 @@ inline auto enum_variety_cosine() -> test_cases_t ivf_pq_inputs y(x); if (y.min_recall.has_value()) { if (y.search_params.lut_dtype == CUDA_R_8U) { - y.search_params.lut_dtype = CUDA_R_16F; - // InnerProduct score is signed, + // CosineExpanded score is signed, // thus we're forced to used signed 8-bit representation, // thus we have one bit less precision - y.min_recall = y.min_recall.value() * 0.90; + y.min_recall = y.min_recall.value() * 0.70; } else { // In other cases it seems to perform a little bit better, still worse than L2 y.min_recall = y.min_recall.value() * 0.94; } } y.index_params.metric = distance::DistanceType::CosineExpanded; + y.index_params.codebook_kind = cuvs::neighbors::ivf_pq::codebook_gen::PER_SUBSPACE; return y; }); } diff --git a/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu index 64c338906..cdc6c1b7e 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu @@ -25,9 +25,9 @@ TEST_BUILD_HOST_INPUT_SEARCH(f32_f32_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_f32_i64) TEST_BUILD_EXTEND_SEARCH(f32_f32_i64) TEST_BUILD_SERIALIZE_SEARCH(f32_f32_i64) -INSTANTIATE(f32_f32_i64, enum_variety_cosine()); +INSTANTIATE(f32_f32_i64, defaults() + small_dims() + big_dims_moderate_lut()); TEST_BUILD_SEARCH(f32_f32_i64_filter) -INSTANTIATE(f32_f32_i64_filter, enum_variety_cosine()); +INSTANTIATE(f32_f32_i64_filter, defaults() + small_dims() + big_dims_moderate_lut()); } // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu index 5c4154bc0..80b0e2ccb 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu @@ -25,8 +25,8 @@ TEST_BUILD_SEARCH(f32_i08_i64) TEST_BUILD_HOST_INPUT_SEARCH(f32_i08_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_i08_i64) TEST_BUILD_SERIALIZE_SEARCH(f32_i08_i64) -INSTANTIATE(f32_i08_i64, enum_variety_cosine()); +INSTANTIATE(f32_i08_i64, defaults() + big_dims() + var_k()); TEST_BUILD_SEARCH(f32_i08_i64_filter) -INSTANTIATE(f32_i08_i64_filter, enum_variety_cosine()); +INSTANTIATE(f32_i08_i64_filter, defaults() + big_dims() + var_k()); } // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu index 3b38fe42b..0216a1e80 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu @@ -25,8 +25,8 @@ TEST_BUILD_SEARCH(f32_u08_i64) TEST_BUILD_HOST_INPUT_SEARCH(f32_u08_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_u08_i64) TEST_BUILD_EXTEND_SEARCH(f32_u08_i64) -INSTANTIATE(f32_u08_i64, small_dims_per_cluster() + enum_variety_cosine()); +INSTANTIATE(f32_u08_i64, small_dims_per_cluster() + enum_variety()); TEST_BUILD_SEARCH(f32_u08_i64_filter) -INSTANTIATE(f32_u08_i64_filter, small_dims_per_cluster() + enum_variety_cosine()); +INSTANTIATE(f32_u08_i64_filter, small_dims_per_cluster() + enum_variety()); } // namespace cuvs::neighbors::ivf_pq From 3ab1c7fb0b19183eb9af042f9da14484919a9d3e Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 25 Sep 2024 14:45:09 -0700 Subject: [PATCH 10/24] style --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 3 ++- cpp/test/neighbors/ann_ivf_pq.cuh | 10 +++------- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index ae215bddc..b3aeb774a 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -1726,7 +1726,8 @@ auto build(raft::resources const& handle, << (int)params.pq_dim << std::endl; RAFT_EXPECTS(n_rows > 0 && dim > 0, "empty dataset"); RAFT_EXPECTS(n_rows >= params.n_lists, "number of rows can't be less than n_lists"); - if (params.metric == cuvs::distance::DistanceType::CosineExpanded && params.codebook_kind == codebook_gen::PER_CLUSTER) { + if (params.metric == cuvs::distance::DistanceType::CosineExpanded && + params.codebook_kind == codebook_gen::PER_CLUSTER) { RAFT_FAIL("CosineExpanded metric only supported for codebook_gen::PER_SUBSPACE"); } diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index 0630e1a8a..faafd3660 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -284,9 +284,7 @@ class ivf_pq_test : public ::testing::TestWithParam { uint32_t n_skip) { // the original data cannot be reconstructed since the dataset was normalized - if (index.metric() == cuvs::distance::DistanceType::CosineExpanded) { - return; - } + if (index.metric() == cuvs::distance::DistanceType::CosineExpanded) { return; } auto& rec_list = index.lists()[label]; auto dim = index.dim(); n_take = std::min(n_take, rec_list->size.load()); @@ -318,9 +316,7 @@ class ivf_pq_test : public ::testing::TestWithParam { auto old_list = index->lists()[label]; auto n_rows = old_list->size.load(); if (n_rows == 0) { return; } - if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { - return; - } + if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { return; } auto vectors_1 = raft::make_device_matrix(handle_, n_rows, index->dim()); auto indices = raft::make_device_vector(handle_, n_rows); @@ -911,7 +907,7 @@ inline auto enum_variety_cosine() -> test_cases_t y.min_recall = y.min_recall.value() * 0.94; } } - y.index_params.metric = distance::DistanceType::CosineExpanded; + y.index_params.metric = distance::DistanceType::CosineExpanded; y.index_params.codebook_kind = cuvs::neighbors::ivf_pq::codebook_gen::PER_SUBSPACE; return y; }); From 6b93282d43498d11fcac1a69764e93ab11362139 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Thu, 26 Sep 2024 09:33:34 -0700 Subject: [PATCH 11/24] use raft::linalg::map --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 23 +++++++---------------- 1 file changed, 7 insertions(+), 16 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index b3aeb774a..f4a272c5b 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -257,7 +257,6 @@ void set_centers(raft::resources const& handle, index* index, const float* raft::linalg::L2Norm, true, stream); - RAFT_CUDA_TRY(cudaMemcpy2DAsync(index->centers().data_handle() + index->dim(), sizeof(float) * index->dim_ext(), center_norms.data(), @@ -1584,13 +1583,7 @@ void extend(raft::resources const& handle, if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { auto float_vec_batch = raft::make_device_matrix(handle, batch.size(), index->dim()); - raft::linalg::map_offset( - handle, - raft::make_device_vector_view(batch.data(), - batch.size() * index->dim()), - raft::make_device_vector_view(float_vec_batch.data_handle(), - float_vec_batch.size()), - [=] __device__(internal_extents_t idx, T i) { return utils::mapping{}(i); }); + raft::linalg::map(handle, float_vec_batch.view(), utils::mapping{}, batch_data_view); raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_vec_batch.view()), float_vec_batch.view(), @@ -1662,13 +1655,11 @@ void extend(raft::resources const& handle, const auto& idx_batch = *idx_batches++; auto float_vec_batch = raft::make_device_matrix(handle, vec_batch.size(), index->dim()); - raft::linalg::map_offset( - handle, - raft::make_device_vector_view(vec_batch.data(), - vec_batch.size() * index->dim()), - raft::make_device_vector_view(float_vec_batch.data_handle(), - vec_batch.size() * index->dim()), - [=] __device__(internal_extents_t idx, T i) { return utils::mapping{}(i); }); + raft::linalg::map(handle, + float_vec_batch.view(), + utils::mapping{}, + raft::make_device_matrix_view( + vec_batch.data(), vec_batch.size(), index->dim())); if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_vec_batch.view()), @@ -1677,7 +1668,7 @@ void extend(raft::resources const& handle, } process_and_fill_codes(handle, *index, - float_vec_batch.data_handle(), + vec_batch.data(), new_indices != nullptr ? std::variant(idx_batch.data()) : std::variant(IdxT(idx_batch.offset())), From bd41e20f409ad00d42975a16eb688c5c80462713 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Thu, 26 Sep 2024 09:57:25 -0700 Subject: [PATCH 12/24] fix ci --- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index b118d52ec..608d5256d 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -716,13 +716,14 @@ inline void search(raft::resources const& handle, rot_queries.data(), index.rot_dim(), stream); - - auto rot_queries_view = raft::make_device_matrix_view( - rot_queries.data(), max_queries, index.rot_dim()); - raft::linalg::row_normalize(handle, - raft::make_const_mdspan(rot_queries_view), - rot_queries_view, - raft::linalg::NormType::L2Norm); + if (index.metric() == distance::DistanceType::CosineExpanded) { + auto rot_queries_view = raft::make_device_matrix_view( + rot_queries.data(), max_queries, index.rot_dim()); + raft::linalg::row_normalize(handle, + raft::make_const_mdspan(rot_queries_view), + rot_queries_view, + raft::linalg::NormType::L2Norm); + } for (uint32_t offset_b = 0; offset_b < queries_batch; offset_b += max_batch_size) { uint32_t batch_size = min(max_batch_size, queries_batch - offset_b); /* The distance calculation is done in the rotated/transformed space; From ca240ac97f0d91b06e3e0ac204cc9aeea40db587 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Thu, 26 Sep 2024 10:15:08 -0700 Subject: [PATCH 13/24] update ivf-flat interleaved scan --- cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh index a4f769741..aca74132c 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh @@ -1205,7 +1205,7 @@ void launch_with_fixed_consts(cuvs::distance::DistanceType metric, Args&&... arg IvfSampleFilterT, inner_prod_dist>( {}, - raft::compose_op(raft::add_const_op{1.0f}, raft::mul_const_op{-1.0f}), + raft::mul_const_op{-1.0f}, std::forward(args)...); // NB: update the description of `knn::ivf_flat::build` when adding here a new metric. default: RAFT_FAIL("The chosen distance metric is not supported (%d)", int(metric)); From 69d1edfb3af3b2345adf643b2ee8db1df0409785 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Thu, 26 Sep 2024 10:18:10 -0700 Subject: [PATCH 14/24] style --- cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh index aca74132c..c5ca1e829 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh @@ -1204,9 +1204,7 @@ void launch_with_fixed_consts(cuvs::distance::DistanceType metric, Args&&... arg IdxT, IvfSampleFilterT, inner_prod_dist>( - {}, - raft::mul_const_op{-1.0f}, - std::forward(args)...); + {}, raft::mul_const_op{-1.0f}, std::forward(args)...); // NB: update the description of `knn::ivf_flat::build` when adding here a new metric. default: RAFT_FAIL("The chosen distance metric is not supported (%d)", int(metric)); } From 28a48d153603acd8ebdef95d08ca3c77da20ec52 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Mon, 30 Sep 2024 13:12:57 -0700 Subject: [PATCH 15/24] rm bug --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index f4a272c5b..cea76cb37 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -1668,7 +1668,7 @@ void extend(raft::resources const& handle, } process_and_fill_codes(handle, *index, - vec_batch.data(), + float_vec_batch.data_handle(), new_indices != nullptr ? std::variant(idx_batch.data()) : std::variant(IdxT(idx_batch.offset())), From 75209a2cd75ad8fe925d7c11ad89ebae6c2bf917 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Mon, 30 Sep 2024 13:38:49 -0700 Subject: [PATCH 16/24] use device_memory mr --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index cea76cb37..5c5f71f7a 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -1581,8 +1581,10 @@ void extend(raft::resources const& handle, cuvs::cluster::kmeans::balanced_params kmeans_params; if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { - auto float_vec_batch = - raft::make_device_matrix(handle, batch.size(), index->dim()); + auto float_vec_batch = raft::make_device_mdarray( + handle, + device_memory, + raft::make_extents(batch.size(), index->dim())); raft::linalg::map(handle, float_vec_batch.view(), utils::mapping{}, batch_data_view); raft::linalg::row_normalize(handle, raft::make_const_mdspan(float_vec_batch.view()), @@ -1653,8 +1655,10 @@ void extend(raft::resources const& handle, vec_batches.prefetch_next_batch(); for (const auto& vec_batch : vec_batches) { const auto& idx_batch = *idx_batches++; - auto float_vec_batch = - raft::make_device_matrix(handle, vec_batch.size(), index->dim()); + auto float_vec_batch = raft::make_device_mdarray( + handle, + device_memory, + raft::make_extents(vec_batch.size(), index->dim())); raft::linalg::map(handle, float_vec_batch.view(), utils::mapping{}, From fca5d94fa0d940dee2783fe127b9863d88729be9 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Mon, 30 Sep 2024 22:12:35 -0700 Subject: [PATCH 17/24] update postprocess --- cpp/src/neighbors/ivf_common.cuh | 9 +-------- .../neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh | 5 +++-- .../ivf_pq/ivf_pq_compute_similarity_impl.cuh | 3 +++ cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 10 ++++++++-- 4 files changed, 15 insertions(+), 12 deletions(-) diff --git a/cpp/src/neighbors/ivf_common.cuh b/cpp/src/neighbors/ivf_common.cuh index 05035e14d..fb73fb8a9 100644 --- a/cpp/src/neighbors/ivf_common.cuh +++ b/cpp/src/neighbors/ivf_common.cuh @@ -254,14 +254,7 @@ void postprocess_distances(ScoreOutT* out, // [n_queries, topk] raft::linalg::unaryOp(out, in, len, raft::sqrt_op{}, stream); } } break; - case distance::DistanceType::CosineExpanded: { - raft::linalg::unaryOp( - out, - in, - len, - raft::compose_op(raft::add_const_op{1.0}, raft::cast_op{}), - stream); - } break; + case distance::DistanceType::CosineExpanded: case distance::DistanceType::InnerProduct: { float factor = (account_for_max_close ? -1.0 : 1.0) * scaling_factor * scaling_factor; if (factor != 1.0) { diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh index c5ca1e829..5c8eb4e42 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh @@ -1204,8 +1204,9 @@ void launch_with_fixed_consts(cuvs::distance::DistanceType metric, Args&&... arg IdxT, IvfSampleFilterT, inner_prod_dist>( - {}, raft::mul_const_op{-1.0f}, std::forward(args)...); - // NB: update the description of `knn::ivf_flat::build` when adding here a new metric. + {}, + raft::compose_op(raft::add_const_op{1.0f}, raft::mul_const_op{-1.0f}), + std::forward(args)...); // NB: update the description of `knn::ivf_flat::build` when adding here a new metric. default: RAFT_FAIL("The chosen distance metric is not supported (%d)", int(metric)); } } diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh index 7a5a2336b..d7a0cd10f 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh @@ -487,6 +487,9 @@ RAFT_KERNEL compute_similarity_kernel(uint32_t dim, reinterpret_cast(pq_thread_data), lut_scores, early_stop_limit); + if (metric == distance::DistanceType::CosineExpanded) { + score = OutT(1) + score; + } } if constexpr (kManageLocalTopK) { block_topk.add(score, sample_offset + i); diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 608d5256d..0c308a070 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -474,8 +474,14 @@ void ivfpq_search_worker(raft::resources const& handle, num_samples_vector); // Postprocessing - ivf::detail::postprocess_distances( - distances, topk_dists.data(), index.metric(), n_queries, topK, scaling_factor, true, stream); + ivf::detail::postprocess_distances(distances, + topk_dists.data(), + index.metric(), + n_queries, + topK, + scaling_factor, + index.metric() != distance::DistanceType::CosineExpanded, + stream); ivf::detail::postprocess_neighbors(neighbors, neighbors_uint32, index.inds_ptrs().data_handle(), From ecf6dc45f59e7bf2c857d30c7c16e85ef460b00b Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Mon, 30 Sep 2024 22:14:42 -0700 Subject: [PATCH 18/24] style --- cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh | 3 ++- cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh | 4 +--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh index 5c8eb4e42..33520d6fc 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_interleaved_scan.cuh @@ -1206,7 +1206,8 @@ void launch_with_fixed_consts(cuvs::distance::DistanceType metric, Args&&... arg inner_prod_dist>( {}, raft::compose_op(raft::add_const_op{1.0f}, raft::mul_const_op{-1.0f}), - std::forward(args)...); // NB: update the description of `knn::ivf_flat::build` when adding here a new metric. + std::forward(args)...); // NB: update the description of `knn::ivf_flat::build` when + // adding here a new metric. default: RAFT_FAIL("The chosen distance metric is not supported (%d)", int(metric)); } } diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh index d7a0cd10f..7192ca2d5 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_compute_similarity_impl.cuh @@ -487,9 +487,7 @@ RAFT_KERNEL compute_similarity_kernel(uint32_t dim, reinterpret_cast(pq_thread_data), lut_scores, early_stop_limit); - if (metric == distance::DistanceType::CosineExpanded) { - score = OutT(1) + score; - } + if (metric == distance::DistanceType::CosineExpanded) { score = OutT(1) + score; } } if constexpr (kManageLocalTopK) { block_topk.add(score, sample_offset + i); From a943f6a96bdb96fe784304aca28d55b91a054902 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 2 Oct 2024 11:26:59 -0700 Subject: [PATCH 19/24] normalize centroids --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 44 ++++++------------- cpp/test/neighbors/ann_ivf_pq.cuh | 2 +- .../ann_ivf_pq/test_float_int64_t.cu | 8 +++- .../ann_ivf_pq/test_int8_t_int64_t.cu | 8 +++- .../ann_ivf_pq/test_uint8_t_int64_t.cu | 8 +++- 5 files changed, 33 insertions(+), 37 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index 5c5f71f7a..5fb0dd27a 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -32,6 +32,7 @@ #include #include #include +#include #include #include #include @@ -1579,33 +1580,13 @@ void extend(raft::resources const& handle, auto centers_view = raft::make_device_matrix_view( cluster_centers.data(), n_clusters, index->dim()); cuvs::cluster::kmeans::balanced_params kmeans_params; - - if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { - auto float_vec_batch = raft::make_device_mdarray( - handle, - device_memory, - raft::make_extents(batch.size(), index->dim())); - raft::linalg::map(handle, float_vec_batch.view(), utils::mapping{}, batch_data_view); - raft::linalg::row_normalize(handle, - raft::make_const_mdspan(float_vec_batch.view()), - float_vec_batch.view(), - raft::linalg::NormType::L2Norm); - - kmeans_params.metric = distance::DistanceType::InnerProduct; - cuvs::cluster::kmeans_balanced::predict(handle, - kmeans_params, - raft::make_const_mdspan(float_vec_batch.view()), - centers_view, - batch_labels_view); - } else { - kmeans_params.metric = static_cast((int)index->metric()); - cuvs::cluster::kmeans_balanced::predict(handle, - kmeans_params, - batch_data_view, - centers_view, - batch_labels_view, - utils::mapping{}); - } + kmeans_params.metric = static_cast((int)index->metric()); + cuvs::cluster::kmeans_balanced::predict(handle, + kmeans_params, + batch_data_view, + centers_view, + batch_labels_view, + utils::mapping{}); vec_batches.prefetch_next_batch(); // User needs to make sure kernel finishes its work before we overwrite batch in the next // iteration if different streams are used for kernel and copy. @@ -1796,12 +1777,11 @@ auto build(raft::resources const& handle, cluster_centers, index.n_lists(), index.dim()); cuvs::cluster::kmeans::balanced_params kmeans_params; kmeans_params.n_iters = params.kmeans_n_iters; + kmeans_params.metric = static_cast((int)index.metric()); + if (index.metric() == distance::DistanceType::CosineExpanded) { raft::linalg::row_normalize( handle, trainset_const_view, trainset.view(), raft::linalg::NormType::L2Norm); - kmeans_params.metric = distance::DistanceType::InnerProduct; - } else { - kmeans_params.metric = static_cast((int)index.metric()); } cuvs::cluster::kmeans_balanced::fit( handle, kmeans_params, trainset_const_view, centers_view, utils::mapping{}); @@ -1810,6 +1790,10 @@ auto build(raft::resources const& handle, rmm::device_uvector labels(n_rows_train, stream, big_memory_resource); auto centers_const_view = raft::make_device_matrix_view( cluster_centers, index.n_lists(), index.dim()); + if (index.metric() == distance::DistanceType::CosineExpanded) { + raft::linalg::row_normalize( + handle, centers_const_view, centers_view, raft::linalg::NormType::L2Norm); + } auto labels_view = raft::make_device_vector_view(labels.data(), n_rows_train); cuvs::cluster::kmeans_balanced::predict(handle, diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index faafd3660..0140a3196 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -378,7 +378,7 @@ class ivf_pq_test : public ::testing::TestWithParam { cuvs::Compare{})); // Pack a few vectors back to the list. - int row_offset = 9; + int row_offset = 5; int n_vec = 3; ASSERT_TRUE(row_offset + n_vec < n_rows); size_t offset = row_offset * index->pq_dim(); diff --git a/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu index cdc6c1b7e..32176d70f 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu @@ -25,9 +25,13 @@ TEST_BUILD_HOST_INPUT_SEARCH(f32_f32_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_f32_i64) TEST_BUILD_EXTEND_SEARCH(f32_f32_i64) TEST_BUILD_SERIALIZE_SEARCH(f32_f32_i64) -INSTANTIATE(f32_f32_i64, defaults() + small_dims() + big_dims_moderate_lut()); +INSTANTIATE(f32_f32_i64, + defaults() + small_dims() + big_dims_moderate_lut() + enum_variety_l2() + + enum_variety_l2sqrt() + enum_variety_ip() + enum_variety_cosine()); TEST_BUILD_SEARCH(f32_f32_i64_filter) -INSTANTIATE(f32_f32_i64_filter, defaults() + small_dims() + big_dims_moderate_lut()); +INSTANTIATE(f32_f32_i64_filter, + enum_variety_ip() + enum_variety_l2() + enum_variety_l2sqrt() + enum_variety_ip() + + enum_variety_cosine()); } // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu index 80b0e2ccb..0d77b718f 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu @@ -25,8 +25,12 @@ TEST_BUILD_SEARCH(f32_i08_i64) TEST_BUILD_HOST_INPUT_SEARCH(f32_i08_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_i08_i64) TEST_BUILD_SERIALIZE_SEARCH(f32_i08_i64) -INSTANTIATE(f32_i08_i64, defaults() + big_dims() + var_k()); +INSTANTIATE(f32_i08_i64, + defaults() + big_dims() + var_k() + enum_variety_l2() + enum_variety_ip() + + enum_variety_cosine()); TEST_BUILD_SEARCH(f32_i08_i64_filter) -INSTANTIATE(f32_i08_i64_filter, defaults() + big_dims() + var_k()); +INSTANTIATE(f32_i08_i64_filter, + defaults() + big_dims() + var_k() + enum_variety_l2() + enum_variety_ip() + + enum_variety_cosine()); } // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu index 0216a1e80..a9a21ecbc 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu @@ -25,8 +25,12 @@ TEST_BUILD_SEARCH(f32_u08_i64) TEST_BUILD_HOST_INPUT_SEARCH(f32_u08_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_u08_i64) TEST_BUILD_EXTEND_SEARCH(f32_u08_i64) -INSTANTIATE(f32_u08_i64, small_dims_per_cluster() + enum_variety()); +INSTANTIATE(f32_u08_i64, + small_dims_per_cluster() + enum_variety() + enum_variety_l2() + enum_variety_l2sqrt() + + enum_variety_ip() + enum_variety_cosine()); TEST_BUILD_SEARCH(f32_u08_i64_filter) -INSTANTIATE(f32_u08_i64_filter, small_dims_per_cluster() + enum_variety()); +INSTANTIATE(f32_u08_i64_filter, + small_dims_per_cluster() + enum_variety() + enum_variety_l2() + enum_variety_l2sqrt() + + enum_variety_ip() + enum_variety_cosine()); } // namespace cuvs::neighbors::ivf_pq From ba660ce09f7aef420d92e14f110e36e21647bbc1 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 2 Oct 2024 12:39:35 -0700 Subject: [PATCH 20/24] allow per_subspace --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 5 ----- cpp/test/neighbors/ann_ivf_pq.cuh | 1 - 2 files changed, 6 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index 5fb0dd27a..1957d3253 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -30,7 +30,6 @@ #include "../../cluster/kmeans_balanced.cuh" #include -#include #include #include #include @@ -1702,10 +1701,6 @@ auto build(raft::resources const& handle, << (int)params.pq_dim << std::endl; RAFT_EXPECTS(n_rows > 0 && dim > 0, "empty dataset"); RAFT_EXPECTS(n_rows >= params.n_lists, "number of rows can't be less than n_lists"); - if (params.metric == cuvs::distance::DistanceType::CosineExpanded && - params.codebook_kind == codebook_gen::PER_CLUSTER) { - RAFT_FAIL("CosineExpanded metric only supported for codebook_gen::PER_SUBSPACE"); - } auto stream = raft::resource::get_cuda_stream(handle); diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index e66f5afba..937bc7c65 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -903,7 +903,6 @@ inline auto enum_variety_cosine() -> test_cases_t } } y.index_params.metric = distance::DistanceType::CosineExpanded; - y.index_params.codebook_kind = cuvs::neighbors::ivf_pq::codebook_gen::PER_SUBSPACE; return y; }); } From bcefbe3b82a61defe4f442868553a401bf3107bf Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 2 Oct 2024 12:45:31 -0700 Subject: [PATCH 21/24] update doc --- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 0a0544614..61fd714de 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -107,6 +107,14 @@ void select_clusters(raft::resources const& handle, This is a negative inner-product distance. We minimize it to find the similar clusters. + NB: qc_distances is NOT used further in ivfpq_search. + + Cosine distance: + `qc_distances[i, j] = - (queries[i], cluster_centers[j])` + + This is a negative inner-product distance. The queries and cluster centers are row normalized. + We minimize it to find the similar clusters. + NB: qc_distances is NOT used further in ivfpq_search. */ float norm_factor; From 496ff5cf874b078155faaf6ca5ba4034c0eb7c61 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Wed, 2 Oct 2024 12:46:43 -0700 Subject: [PATCH 22/24] style --- cpp/test/neighbors/ann_ivf_pq.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index 937bc7c65..a35a0cb88 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -902,7 +902,7 @@ inline auto enum_variety_cosine() -> test_cases_t y.min_recall = y.min_recall.value() * 0.94; } } - y.index_params.metric = distance::DistanceType::CosineExpanded; + y.index_params.metric = distance::DistanceType::CosineExpanded; return y; }); } From 490e6d2653749ba7f17d9142aa397e7b8b4ba46b Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Thu, 3 Oct 2024 11:52:38 -0700 Subject: [PATCH 23/24] only support float --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 33 +++++++++++-------- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 19 ++++------- cpp/test/neighbors/ann_ivf_pq.cuh | 5 ++- .../ann_ivf_pq/test_int8_t_int64_t.cu | 7 ++-- .../ann_ivf_pq/test_uint8_t_int64_t.cu | 4 +-- 5 files changed, 33 insertions(+), 35 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index 1957d3253..4c9867126 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -1469,6 +1469,13 @@ void extend(raft::resources const& handle, std::is_same_v, "Unsupported data type"); + if (index->metric() == distance::DistanceType::CosineExpanded) { + if constexpr (std::is_same_v || std::is_same_v) + RAFT_FAIL( + "CosineExpanded distance metric is currently not supported for uint8_t and int8_t data " + "type"); + } + rmm::device_async_resource_ref device_memory = raft::resource::get_workspace_resource(handle); rmm::device_async_resource_ref large_memory = raft::resource::get_large_workspace_resource(handle); @@ -1635,24 +1642,17 @@ void extend(raft::resources const& handle, vec_batches.prefetch_next_batch(); for (const auto& vec_batch : vec_batches) { const auto& idx_batch = *idx_batches++; - auto float_vec_batch = raft::make_device_mdarray( - handle, - device_memory, - raft::make_extents(vec_batch.size(), index->dim())); - raft::linalg::map(handle, - float_vec_batch.view(), - utils::mapping{}, - raft::make_device_matrix_view( - vec_batch.data(), vec_batch.size(), index->dim())); - if (index->metric() == cuvs::distance::DistanceType::CosineExpanded) { + if (index->metric() == CosineExpanded) { + auto vec_batch_view = raft::make_device_matrix_view( + const_cast(vec_batch.data()), vec_batch.size(), index->dim()); raft::linalg::row_normalize(handle, - raft::make_const_mdspan(float_vec_batch.view()), - float_vec_batch.view(), + raft::make_const_mdspan(vec_batch_view), + vec_batch_view, raft::linalg::NormType::L2Norm); } process_and_fill_codes(handle, *index, - float_vec_batch.data_handle(), + vec_batch.data(), new_indices != nullptr ? std::variant(idx_batch.data()) : std::variant(IdxT(idx_batch.offset())), @@ -1701,6 +1701,13 @@ auto build(raft::resources const& handle, << (int)params.pq_dim << std::endl; RAFT_EXPECTS(n_rows > 0 && dim > 0, "empty dataset"); RAFT_EXPECTS(n_rows >= params.n_lists, "number of rows can't be less than n_lists"); + if (params.metric == distance::DistanceType::CosineExpanded) { + // TODO: support int8_t and uint8_t types (https://github.com/rapidsai/cuvs/issues/389) + if constexpr (std::is_same_v || std::is_same_v) + RAFT_FAIL( + "CosineExpanded distance metric is currently not supported for uint8_t and int8_t data " + "type"); + } auto stream = raft::resource::get_cuda_stream(handle); diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 61fd714de..db8f9fbd3 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -145,22 +145,11 @@ void select_clusters(raft::resources const& handle, gemm_k = dim + 1; RAFT_EXPECTS(gemm_k <= dim_ext, "unexpected gemm_k or dim_ext"); } break; + case cuvs::distance::DistanceType::CosineExpanded: case cuvs::distance::DistanceType::InnerProduct: { alpha = -1.0; beta = 0.0; } break; - case cuvs::distance::DistanceType::CosineExpanded: { - alpha = -1.0; - beta = 0.0; - - auto float_queries_matrix_view = - raft::make_device_matrix_view(float_queries, n_queries, dim_ext); - - raft::linalg::row_normalize(handle, - raft::make_const_mdspan(float_queries_matrix_view), - float_queries_matrix_view, - raft::linalg::NormType::L2Norm); - } break; default: RAFT_FAIL("Unsupported distance type %d.", int(metric)); } rmm::device_uvector qc_distances(n_queries * n_lists, stream, mr); @@ -638,6 +627,12 @@ inline void search(raft::resources const& handle, static_assert(std::is_same_v || std::is_same_v || std::is_same_v || std::is_same_v, "Unsupported element type."); + if (index.metric() == distance::DistanceType::CosineExpanded) { + if constexpr (std::is_same_v || std::is_same_v) + RAFT_FAIL( + "CosineExpanded distance metric is currently not supported for uint8_t and int8_t data " + "type"); + } raft::common::nvtx::range fun_scope( "ivf_pq::search(n_queries = %u, n_probes = %u, k = %u, dim = %zu)", n_queries, diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index a35a0cb88..fd4e330db 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -893,9 +893,8 @@ inline auto enum_variety_cosine() -> test_cases_t ivf_pq_inputs y(x); if (y.min_recall.has_value()) { if (y.search_params.lut_dtype == CUDA_R_8U) { - // CosineExpanded score is signed, - // thus we're forced to used signed 8-bit representation, - // thus we have one bit less precision + // TODO: Increase this recall threshold for 8 bit lut + // (https://github.com/rapidsai/cuvs/issues/390) y.min_recall = y.min_recall.value() * 0.70; } else { // In other cases it seems to perform a little bit better, still worse than L2 diff --git a/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu index 0d77b718f..c9e5d4f01 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_int8_t_int64_t.cu @@ -25,12 +25,9 @@ TEST_BUILD_SEARCH(f32_i08_i64) TEST_BUILD_HOST_INPUT_SEARCH(f32_i08_i64) TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_i08_i64) TEST_BUILD_SERIALIZE_SEARCH(f32_i08_i64) -INSTANTIATE(f32_i08_i64, - defaults() + big_dims() + var_k() + enum_variety_l2() + enum_variety_ip() + - enum_variety_cosine()); +INSTANTIATE(f32_i08_i64, defaults() + big_dims() + var_k() + enum_variety_l2() + enum_variety_ip()); TEST_BUILD_SEARCH(f32_i08_i64_filter) INSTANTIATE(f32_i08_i64_filter, - defaults() + big_dims() + var_k() + enum_variety_l2() + enum_variety_ip() + - enum_variety_cosine()); + defaults() + big_dims() + var_k() + enum_variety_l2() + enum_variety_ip()); } // namespace cuvs::neighbors::ivf_pq diff --git a/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu index a9a21ecbc..6e0732227 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_uint8_t_int64_t.cu @@ -27,10 +27,10 @@ TEST_BUILD_HOST_INPUT_OVERLAP_SEARCH(f32_u08_i64) TEST_BUILD_EXTEND_SEARCH(f32_u08_i64) INSTANTIATE(f32_u08_i64, small_dims_per_cluster() + enum_variety() + enum_variety_l2() + enum_variety_l2sqrt() + - enum_variety_ip() + enum_variety_cosine()); + enum_variety_ip()); TEST_BUILD_SEARCH(f32_u08_i64_filter) INSTANTIATE(f32_u08_i64_filter, small_dims_per_cluster() + enum_variety() + enum_variety_l2() + enum_variety_l2sqrt() + - enum_variety_ip() + enum_variety_cosine()); + enum_variety_ip()); } // namespace cuvs::neighbors::ivf_pq From eac9bba7e90de7926ca37fafcc350e9cb2dfe8b5 Mon Sep 17 00:00:00 2001 From: Tarang Jain Date: Thu, 3 Oct 2024 12:23:51 -0700 Subject: [PATCH 24/24] run remaining float tests --- cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu b/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu index 32176d70f..834fdb3d0 100644 --- a/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu +++ b/cpp/test/neighbors/ann_ivf_pq/test_float_int64_t.cu @@ -31,7 +31,7 @@ INSTANTIATE(f32_f32_i64, TEST_BUILD_SEARCH(f32_f32_i64_filter) INSTANTIATE(f32_f32_i64_filter, - enum_variety_ip() + enum_variety_l2() + enum_variety_l2sqrt() + enum_variety_ip() + - enum_variety_cosine()); + defaults() + small_dims() + big_dims_moderate_lut() + enum_variety_l2() + + enum_variety_l2sqrt() + enum_variety_ip() + enum_variety_cosine()); } // namespace cuvs::neighbors::ivf_pq