From 59bb8ca9afeaddc0b7dce234f2d49df000254394 Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Tue, 5 Nov 2024 13:45:25 +0000 Subject: [PATCH 01/16] . --- cpp/include/cuvs/neighbors/cagra.h | 16 ++++ cpp/src/neighbors/cagra_c.cpp | 77 ++++++++++++++++++- cpp/test/neighbors/ann_cagra_c.cu | 117 +++++++++++++++++++++++++++++ 3 files changed, 207 insertions(+), 3 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 14331ebbca..b629e5c1f5 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -392,6 +392,22 @@ cuvsError_t cuvsCagraSearch(cuvsResources_t res, DLManagedTensor* queries, DLManagedTensor* neighbors, DLManagedTensor* distances); +/* + * @param[in] res cuvsResources_t opaque C handle + * @param[in] params cuvsCagraSearchParams_t used to search CAGRA index + * @param[in] index cuvsCagraIndex which has been returned by `cuvsCagraBuild` + * @param[in] queries DLManagedTensor* queries dataset to search + * @param[in] filter Filter + * @param[out] neighbors DLManagedTensor* output `k` neighbors for queries + * @param[out] distances DLManagedTensor* output `k` distances for queries + */ +cuvsError_t cuvsCagraFilteredSearch(cuvsResources_t res, + cuvsCagraSearchParams_t params, + cuvsCagraIndex_t index, + DLManagedTensor* queries, + DLManagedTensor* neighbors, + DLManagedTensor* distances, + DLManagedTensor* filter); /** * @} diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 6985ff0948..bb178de7c9 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -29,6 +29,8 @@ #include #include +#include + namespace { template @@ -89,7 +91,8 @@ void _search(cuvsResources_t res, cuvsCagraIndex index, DLManagedTensor* queries_tensor, DLManagedTensor* neighbors_tensor, - DLManagedTensor* distances_tensor) + DLManagedTensor* distances_tensor, + std::optional removed_indices_tensor = std::nullopt) { auto res_ptr = reinterpret_cast(res); auto index_ptr = reinterpret_cast*>(index.addr); @@ -115,8 +118,26 @@ void _search(cuvsResources_t res, auto queries_mds = cuvs::core::from_dlpack(queries_tensor); auto neighbors_mds = cuvs::core::from_dlpack(neighbors_tensor); auto distances_mds = cuvs::core::from_dlpack(distances_tensor); - cuvs::neighbors::cagra::search( - *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds); + + if (removed_indices_tensor.has_value()) { + using filter_mdspan_type = raft::device_vector_view; + auto removed_indices = + cuvs::core::from_dlpack(removed_indices_tensor.value()); + cuvs::core::bitset removed_indices_bitset( + *res_ptr, removed_indices, index_ptr->dataset().extent(0)); + auto bitset_filter_obj = + cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset.view()); + cuvs::neighbors::cagra::search(*res_ptr, + search_params, + *index_ptr, + queries_mds, + neighbors_mds, + distances_mds, + bitset_filter_obj); + } else { + cuvs::neighbors::cagra::search( + *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds); + } } template @@ -247,6 +268,56 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res, }); } +extern "C" cuvsError_t cuvsCagraFilteredSearch(cuvsResources_t res, + cuvsCagraSearchParams_t params, + cuvsCagraIndex_t index_c_ptr, + DLManagedTensor* queries_tensor, + DLManagedTensor* neighbors_tensor, + DLManagedTensor* distances_tensor, + DLManagedTensor* filter_tensor) +{ + return cuvs::core::translate_exceptions([=] { + auto queries = queries_tensor->dl_tensor; + auto neighbors = neighbors_tensor->dl_tensor; + auto distances = distances_tensor->dl_tensor; + auto filter = filter_tensor->dl_tensor; + + RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(queries), + "queries should have device compatible memory"); + RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(neighbors), + "neighbors should have device compatible memory"); + RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(distances), + "distances should have device compatible memory"); + RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(filter), + "filter should have device compatible memory"); + + RAFT_EXPECTS(neighbors.dtype.code == kDLUInt && neighbors.dtype.bits == 32, + "neighbors should be of type uint32_t"); + RAFT_EXPECTS(distances.dtype.code == kDLFloat && neighbors.dtype.bits == 32, + "distances should be of type float32"); + RAFT_EXPECTS(filter.dtype.code == kDLInt && filter.dtype.bits == 64, + "filter should be of type int64_t"); + + auto index = *index_c_ptr; + RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between index and queries"); + + if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { + _search( + res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter_tensor); + } else if (queries.dtype.code == kDLInt && queries.dtype.bits == 8) { + _search( + res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter_tensor); + } else if (queries.dtype.code == kDLUInt && queries.dtype.bits == 8) { + _search( + res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter_tensor); + } else { + RAFT_FAIL("Unsupported queries DLtensor dtype: %d and bits: %d", + queries.dtype.code, + queries.dtype.bits); + } + }); +} + extern "C" cuvsError_t cuvsCagraIndexParamsCreate(cuvsCagraIndexParams_t* params) { return cuvs::core::translate_exceptions([=] { diff --git a/cpp/test/neighbors/ann_cagra_c.cu b/cpp/test/neighbors/ann_cagra_c.cu index 599d2d8428..9330cf5731 100644 --- a/cpp/test/neighbors/ann_cagra_c.cu +++ b/cpp/test/neighbors/ann_cagra_c.cu @@ -34,6 +34,8 @@ float queries[4][2] = {{0.48216683, 0.0428398}, {0.51260436, 0.2643005}, {0.05198065, 0.5789965}}; +int64_t filter[2] = {1, 2}; + uint32_t neighbors_exp[4] = {3, 0, 3, 1}; float distances_exp[4] = {0.03878258, 0.12472608, 0.04776672, 0.15224178}; @@ -126,3 +128,118 @@ TEST(CagraC, BuildSearch) cuvsCagraIndexDestroy(index); cuvsResourcesDestroy(res); } + +TEST(CagraC, BuildSearchFiltered) +{ + // create cuvsResources_t + cuvsResources_t res; + cuvsResourcesCreate(&res); + cudaStream_t stream; + cuvsStreamGet(res, &stream); + + // create dataset DLTensor + DLManagedTensor dataset_tensor; + dataset_tensor.dl_tensor.data = dataset; + dataset_tensor.dl_tensor.device.device_type = kDLCPU; + dataset_tensor.dl_tensor.ndim = 2; + dataset_tensor.dl_tensor.dtype.code = kDLFloat; + dataset_tensor.dl_tensor.dtype.bits = 32; + dataset_tensor.dl_tensor.dtype.lanes = 1; + int64_t dataset_shape[2] = {4, 2}; + dataset_tensor.dl_tensor.shape = dataset_shape; + dataset_tensor.dl_tensor.strides = nullptr; + + // create index + cuvsCagraIndex_t index; + cuvsCagraIndexCreate(&index); + + // build index + cuvsCagraIndexParams_t build_params; + cuvsCagraIndexParamsCreate(&build_params); + cuvsCagraBuild(res, build_params, &dataset_tensor, index); + + // create queries DLTensor + rmm::device_uvector queries_d(4 * 2, stream); + raft::copy(queries_d.data(), (float*)queries, 4 * 2, stream); + + DLManagedTensor queries_tensor; + queries_tensor.dl_tensor.data = queries_d.data(); + queries_tensor.dl_tensor.device.device_type = kDLCUDA; + queries_tensor.dl_tensor.ndim = 2; + queries_tensor.dl_tensor.dtype.code = kDLFloat; + queries_tensor.dl_tensor.dtype.bits = 32; + queries_tensor.dl_tensor.dtype.lanes = 1; + int64_t queries_shape[2] = {4, 2}; + queries_tensor.dl_tensor.shape = queries_shape; + queries_tensor.dl_tensor.strides = nullptr; + + // create filter DLTensor + rmm::device_uvector filter_d(2, stream); + raft::copy(filter_d.data(), (int64_t*)filter, 2, stream); + + DLManagedTensor filter_tensor; + filter_tensor.dl_tensor.data = queries_d.data(); + filter_tensor.dl_tensor.device.device_type = kDLCUDA; + filter_tensor.dl_tensor.ndim = 1; + filter_tensor.dl_tensor.dtype.code = kDLInt; + filter_tensor.dl_tensor.dtype.bits = 64; + filter_tensor.dl_tensor.dtype.lanes = 1; + int64_t filter_shape[1] = {2}; + filter_tensor.dl_tensor.shape = filter_shape; + filter_tensor.dl_tensor.strides = nullptr; + + // create neighbors DLTensor + rmm::device_uvector neighbors_d(4, stream); + + DLManagedTensor neighbors_tensor; + neighbors_tensor.dl_tensor.data = neighbors_d.data(); + neighbors_tensor.dl_tensor.device.device_type = kDLCUDA; + neighbors_tensor.dl_tensor.ndim = 2; + neighbors_tensor.dl_tensor.dtype.code = kDLUInt; + neighbors_tensor.dl_tensor.dtype.bits = 32; + neighbors_tensor.dl_tensor.dtype.lanes = 1; + int64_t neighbors_shape[2] = {4, 1}; + neighbors_tensor.dl_tensor.shape = neighbors_shape; + neighbors_tensor.dl_tensor.strides = nullptr; + + // create distances DLTensor + rmm::device_uvector distances_d(4, stream); + + DLManagedTensor distances_tensor; + distances_tensor.dl_tensor.data = distances_d.data(); + distances_tensor.dl_tensor.device.device_type = kDLCUDA; + distances_tensor.dl_tensor.ndim = 2; + distances_tensor.dl_tensor.dtype.code = kDLFloat; + distances_tensor.dl_tensor.dtype.bits = 32; + distances_tensor.dl_tensor.dtype.lanes = 1; + int64_t distances_shape[2] = {4, 1}; + distances_tensor.dl_tensor.shape = distances_shape; + distances_tensor.dl_tensor.strides = nullptr; + + // search index + cuvsCagraSearchParams_t search_params; + cuvsCagraSearchParamsCreate(&search_params); + auto e = cuvsCagraFilteredSearch(res, + search_params, + index, + &queries_tensor, + &neighbors_tensor, + &distances_tensor, + &filter_tensor); + + // if (e != cuvsError_t::CUVS_SUCCESS) { + // std::cout << "Error: " << cuvsGetLastErrorText() << std::endl; + // } + + // verify output + ASSERT_TRUE( + cuvs::devArrMatchHost(neighbors_exp, neighbors_d.data(), 4, cuvs::Compare())); + ASSERT_TRUE(cuvs::devArrMatchHost( + distances_exp, distances_d.data(), 4, cuvs::CompareApprox(0.001f))); + + // de-allocate index and res + cuvsCagraSearchParamsDestroy(search_params); + cuvsCagraIndexParamsDestroy(build_params); + cuvsCagraIndexDestroy(index); + cuvsResourcesDestroy(res); +} From f37c58ee677cece3204ba2deb8cd0a95e1bef6ce Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Fri, 8 Nov 2024 16:07:51 +0000 Subject: [PATCH 02/16] fix test --- cpp/test/neighbors/ann_cagra_c.cu | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/cpp/test/neighbors/ann_cagra_c.cu b/cpp/test/neighbors/ann_cagra_c.cu index 9330cf5731..6e5008d581 100644 --- a/cpp/test/neighbors/ann_cagra_c.cu +++ b/cpp/test/neighbors/ann_cagra_c.cu @@ -39,6 +39,9 @@ int64_t filter[2] = {1, 2}; uint32_t neighbors_exp[4] = {3, 0, 3, 1}; float distances_exp[4] = {0.03878258, 0.12472608, 0.04776672, 0.15224178}; +uint32_t neighbors_exp_filtered[4] = {3, 0, 3, 0}; +float distances_exp_filtered[4] = {0.03878258, 0.12472608, 0.04776672, 0.59063464}; + TEST(CagraC, BuildSearch) { // create cuvsResources_t @@ -178,7 +181,7 @@ TEST(CagraC, BuildSearchFiltered) raft::copy(filter_d.data(), (int64_t*)filter, 2, stream); DLManagedTensor filter_tensor; - filter_tensor.dl_tensor.data = queries_d.data(); + filter_tensor.dl_tensor.data = filter_d.data(); filter_tensor.dl_tensor.device.device_type = kDLCUDA; filter_tensor.dl_tensor.ndim = 1; filter_tensor.dl_tensor.dtype.code = kDLInt; @@ -232,10 +235,10 @@ TEST(CagraC, BuildSearchFiltered) // } // verify output - ASSERT_TRUE( - cuvs::devArrMatchHost(neighbors_exp, neighbors_d.data(), 4, cuvs::Compare())); ASSERT_TRUE(cuvs::devArrMatchHost( - distances_exp, distances_d.data(), 4, cuvs::CompareApprox(0.001f))); + neighbors_exp_filtered, neighbors_d.data(), 4, cuvs::Compare())); + ASSERT_TRUE(cuvs::devArrMatchHost( + distances_exp_filtered, distances_d.data(), 4, cuvs::CompareApprox(0.001f))); // de-allocate index and res cuvsCagraSearchParamsDestroy(search_params); From 4fd78ca86bffd96ea244796b04204a317e9e7ae9 Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Fri, 8 Nov 2024 16:12:39 +0000 Subject: [PATCH 03/16] cleanup --- cpp/src/neighbors/cagra_c.cpp | 2 -- cpp/test/neighbors/ann_cagra_c.cu | 4 ---- 2 files changed, 6 deletions(-) diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index bb178de7c9..c3631d40ad 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -29,8 +29,6 @@ #include #include -#include - namespace { template diff --git a/cpp/test/neighbors/ann_cagra_c.cu b/cpp/test/neighbors/ann_cagra_c.cu index 6e5008d581..2b6e7e5e2a 100644 --- a/cpp/test/neighbors/ann_cagra_c.cu +++ b/cpp/test/neighbors/ann_cagra_c.cu @@ -230,10 +230,6 @@ TEST(CagraC, BuildSearchFiltered) &distances_tensor, &filter_tensor); - // if (e != cuvsError_t::CUVS_SUCCESS) { - // std::cout << "Error: " << cuvsGetLastErrorText() << std::endl; - // } - // verify output ASSERT_TRUE(cuvs::devArrMatchHost( neighbors_exp_filtered, neighbors_d.data(), 4, cuvs::Compare())); From fad44c2a52318efbcf93e688a2a8a042e1601eef Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Wed, 27 Nov 2024 12:26:53 +0000 Subject: [PATCH 04/16] address comments --- cpp/include/cuvs/neighbors/cagra.h | 26 ++--------- cpp/src/neighbors/cagra_c.cpp | 74 ++++++------------------------ cpp/test/neighbors/ann_cagra_c.cu | 51 ++++++++++---------- 3 files changed, 48 insertions(+), 103 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index b629e5c1f5..6abe43e381 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -17,6 +17,7 @@ #pragma once #include +#include #include #include #include @@ -385,33 +386,16 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res, * @param[in] queries DLManagedTensor* queries dataset to search * @param[out] neighbors DLManagedTensor* output `k` neighbors for queries * @param[out] distances DLManagedTensor* output `k` distances for queries + * @param[in] prefilter cuvsFilter input prefilter that can be used + to filter queries and neighbors based on the given bitset. */ cuvsError_t cuvsCagraSearch(cuvsResources_t res, cuvsCagraSearchParams_t params, cuvsCagraIndex_t index, DLManagedTensor* queries, DLManagedTensor* neighbors, - DLManagedTensor* distances); -/* - * @param[in] res cuvsResources_t opaque C handle - * @param[in] params cuvsCagraSearchParams_t used to search CAGRA index - * @param[in] index cuvsCagraIndex which has been returned by `cuvsCagraBuild` - * @param[in] queries DLManagedTensor* queries dataset to search - * @param[in] filter Filter - * @param[out] neighbors DLManagedTensor* output `k` neighbors for queries - * @param[out] distances DLManagedTensor* output `k` distances for queries - */ -cuvsError_t cuvsCagraFilteredSearch(cuvsResources_t res, - cuvsCagraSearchParams_t params, - cuvsCagraIndex_t index, - DLManagedTensor* queries, - DLManagedTensor* neighbors, - DLManagedTensor* distances, - DLManagedTensor* filter); - -/** - * @} - */ + DLManagedTensor* distances, + cuvsFilter filter); /** * @defgroup cagra_c_serialize CAGRA C-API serialize functions diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index 25408ff012..ade29fa063 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #include @@ -92,7 +93,7 @@ void _search(cuvsResources_t res, DLManagedTensor* queries_tensor, DLManagedTensor* neighbors_tensor, DLManagedTensor* distances_tensor, - std::optional removed_indices_tensor = std::nullopt) + cuvsFilter filter) { auto res_ptr = reinterpret_cast(res); auto index_ptr = reinterpret_cast*>(index.addr); @@ -118,11 +119,13 @@ void _search(cuvsResources_t res, auto queries_mds = cuvs::core::from_dlpack(queries_tensor); auto neighbors_mds = cuvs::core::from_dlpack(neighbors_tensor); auto distances_mds = cuvs::core::from_dlpack(distances_tensor); - - if (removed_indices_tensor.has_value()) { - using filter_mdspan_type = raft::device_vector_view; - auto removed_indices = - cuvs::core::from_dlpack(removed_indices_tensor.value()); + if (filter.type == NO_FILTER) { + cuvs::neighbors::cagra::search( + *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds); + } else if (filter.type == BITSET) { + using filter_mdspan_type = raft::device_vector_view; + auto removed_indices_tensor = reinterpret_cast(filter.addr); + auto removed_indices = cuvs::core::from_dlpack(removed_indices_tensor); cuvs::core::bitset removed_indices_bitset( *res_ptr, removed_indices, index_ptr->dataset().extent(0)); auto bitset_filter_obj = @@ -135,8 +138,7 @@ void _search(cuvsResources_t res, distances_mds, bitset_filter_obj); } else { - cuvs::neighbors::cagra::search( - *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds); + RAFT_FAIL("Unsupported prefilter type: BITMAP"); } } @@ -232,55 +234,13 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res, cuvsCagraIndex_t index_c_ptr, DLManagedTensor* queries_tensor, DLManagedTensor* neighbors_tensor, - DLManagedTensor* distances_tensor) -{ - return cuvs::core::translate_exceptions([=] { - auto queries = queries_tensor->dl_tensor; - auto neighbors = neighbors_tensor->dl_tensor; - auto distances = distances_tensor->dl_tensor; - - RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(queries), - "queries should have device compatible memory"); - RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(neighbors), - "neighbors should have device compatible memory"); - RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(distances), - "distances should have device compatible memory"); - - RAFT_EXPECTS(neighbors.dtype.code == kDLUInt && neighbors.dtype.bits == 32, - "neighbors should be of type uint32_t"); - RAFT_EXPECTS(distances.dtype.code == kDLFloat && neighbors.dtype.bits == 32, - "distances should be of type float32"); - - auto index = *index_c_ptr; - RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between index and queries"); - - if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { - _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); - } else if (queries.dtype.code == kDLInt && queries.dtype.bits == 8) { - _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); - } else if (queries.dtype.code == kDLUInt && queries.dtype.bits == 8) { - _search(res, *params, index, queries_tensor, neighbors_tensor, distances_tensor); - } else { - RAFT_FAIL("Unsupported queries DLtensor dtype: %d and bits: %d", - queries.dtype.code, - queries.dtype.bits); - } - }); -} - -extern "C" cuvsError_t cuvsCagraFilteredSearch(cuvsResources_t res, - cuvsCagraSearchParams_t params, - cuvsCagraIndex_t index_c_ptr, - DLManagedTensor* queries_tensor, - DLManagedTensor* neighbors_tensor, - DLManagedTensor* distances_tensor, - DLManagedTensor* filter_tensor) + DLManagedTensor* distances_tensor, + cuvsFilter filter) { return cuvs::core::translate_exceptions([=] { auto queries = queries_tensor->dl_tensor; auto neighbors = neighbors_tensor->dl_tensor; auto distances = distances_tensor->dl_tensor; - auto filter = filter_tensor->dl_tensor; RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(queries), "queries should have device compatible memory"); @@ -288,28 +248,24 @@ extern "C" cuvsError_t cuvsCagraFilteredSearch(cuvsResources_t res, "neighbors should have device compatible memory"); RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(distances), "distances should have device compatible memory"); - RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(filter), - "filter should have device compatible memory"); RAFT_EXPECTS(neighbors.dtype.code == kDLUInt && neighbors.dtype.bits == 32, "neighbors should be of type uint32_t"); RAFT_EXPECTS(distances.dtype.code == kDLFloat && neighbors.dtype.bits == 32, "distances should be of type float32"); - RAFT_EXPECTS(filter.dtype.code == kDLInt && filter.dtype.bits == 64, - "filter should be of type int64_t"); auto index = *index_c_ptr; RAFT_EXPECTS(queries.dtype.code == index.dtype.code, "type mismatch between index and queries"); if (queries.dtype.code == kDLFloat && queries.dtype.bits == 32) { _search( - res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter_tensor); + res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter); } else if (queries.dtype.code == kDLInt && queries.dtype.bits == 8) { _search( - res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter_tensor); + res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter); } else if (queries.dtype.code == kDLUInt && queries.dtype.bits == 8) { _search( - res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter_tensor); + res, *params, index, queries_tensor, neighbors_tensor, distances_tensor, filter); } else { RAFT_FAIL("Unsupported queries DLtensor dtype: %d and bits: %d", queries.dtype.code, diff --git a/cpp/test/neighbors/ann_cagra_c.cu b/cpp/test/neighbors/ann_cagra_c.cu index 2b6e7e5e2a..1a38d7b5ef 100644 --- a/cpp/test/neighbors/ann_cagra_c.cu +++ b/cpp/test/neighbors/ann_cagra_c.cu @@ -114,10 +114,15 @@ TEST(CagraC, BuildSearch) distances_tensor.dl_tensor.shape = distances_shape; distances_tensor.dl_tensor.strides = nullptr; + cuvsFilter prefilter; + prefilter.type = NO_FILTER; + prefilter.addr = (uintptr_t)NULL; + // search index cuvsCagraSearchParams_t search_params; cuvsCagraSearchParamsCreate(&search_params); - cuvsCagraSearch(res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor); + cuvsCagraSearch( + res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, prefilter); // verify output ASSERT_TRUE( @@ -176,21 +181,6 @@ TEST(CagraC, BuildSearchFiltered) queries_tensor.dl_tensor.shape = queries_shape; queries_tensor.dl_tensor.strides = nullptr; - // create filter DLTensor - rmm::device_uvector filter_d(2, stream); - raft::copy(filter_d.data(), (int64_t*)filter, 2, stream); - - DLManagedTensor filter_tensor; - filter_tensor.dl_tensor.data = filter_d.data(); - filter_tensor.dl_tensor.device.device_type = kDLCUDA; - filter_tensor.dl_tensor.ndim = 1; - filter_tensor.dl_tensor.dtype.code = kDLInt; - filter_tensor.dl_tensor.dtype.bits = 64; - filter_tensor.dl_tensor.dtype.lanes = 1; - int64_t filter_shape[1] = {2}; - filter_tensor.dl_tensor.shape = filter_shape; - filter_tensor.dl_tensor.strides = nullptr; - // create neighbors DLTensor rmm::device_uvector neighbors_d(4, stream); @@ -219,16 +209,31 @@ TEST(CagraC, BuildSearchFiltered) distances_tensor.dl_tensor.shape = distances_shape; distances_tensor.dl_tensor.strides = nullptr; + // create filter DLTensor + rmm::device_uvector filter_d(2, stream); + raft::copy(filter_d.data(), (int64_t*)filter, 2, stream); + + cuvsFilter filter; + + DLManagedTensor filter_tensor; + filter_tensor.dl_tensor.data = filter_d.data(); + filter_tensor.dl_tensor.device.device_type = kDLCUDA; + filter_tensor.dl_tensor.ndim = 1; + filter_tensor.dl_tensor.dtype.code = kDLInt; + filter_tensor.dl_tensor.dtype.bits = 64; + filter_tensor.dl_tensor.dtype.lanes = 1; + int64_t filter_shape[1] = {2}; + filter_tensor.dl_tensor.shape = filter_shape; + filter_tensor.dl_tensor.strides = nullptr; + + filter.type = BITSET; + filter.addr = (uintptr_t)&filter_tensor; + // search index cuvsCagraSearchParams_t search_params; cuvsCagraSearchParamsCreate(&search_params); - auto e = cuvsCagraFilteredSearch(res, - search_params, - index, - &queries_tensor, - &neighbors_tensor, - &distances_tensor, - &filter_tensor); + cuvsCagraSearch( + res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter); // verify output ASSERT_TRUE(cuvs::devArrMatchHost( From 0d0184a9491c46d880809edf852c89ae550d33fe Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Wed, 27 Nov 2024 12:28:38 +0000 Subject: [PATCH 05/16] . --- cpp/include/cuvs/neighbors/cagra.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 6abe43e381..95e4ea2604 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -397,6 +397,10 @@ cuvsError_t cuvsCagraSearch(cuvsResources_t res, DLManagedTensor* distances, cuvsFilter filter); +/** + * @} + */ + /** * @defgroup cagra_c_serialize CAGRA C-API serialize functions * @{ From e3d33a064ddc1f17803798fe4589cc6addc41a32 Mon Sep 17 00:00:00 2001 From: Ajit Mistry <55892788+ajit283@users.noreply.github.com> Date: Fri, 29 Nov 2024 01:21:15 +0100 Subject: [PATCH 06/16] Update cpp/src/neighbors/cagra_c.cpp Co-authored-by: Micka --- cpp/src/neighbors/cagra_c.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index ade29fa063..dc9abe0701 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -123,7 +123,7 @@ void _search(cuvsResources_t res, cuvs::neighbors::cagra::search( *res_ptr, search_params, *index_ptr, queries_mds, neighbors_mds, distances_mds); } else if (filter.type == BITSET) { - using filter_mdspan_type = raft::device_vector_view; + using filter_mdspan_type = raft::device_vector_view; auto removed_indices_tensor = reinterpret_cast(filter.addr); auto removed_indices = cuvs::core::from_dlpack(removed_indices_tensor); cuvs::core::bitset removed_indices_bitset( From 3b178953bd317d1adf97a737903b47f09eec14c8 Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Tue, 3 Dec 2024 18:11:17 +0000 Subject: [PATCH 07/16] fix filter repr. --- cpp/src/neighbors/cagra_c.cpp | 9 ++++----- cpp/test/neighbors/ann_cagra_c.cu | 21 ++++++++++----------- 2 files changed, 14 insertions(+), 16 deletions(-) diff --git a/cpp/src/neighbors/cagra_c.cpp b/cpp/src/neighbors/cagra_c.cpp index dc9abe0701..639f32d8e9 100644 --- a/cpp/src/neighbors/cagra_c.cpp +++ b/cpp/src/neighbors/cagra_c.cpp @@ -126,10 +126,9 @@ void _search(cuvsResources_t res, using filter_mdspan_type = raft::device_vector_view; auto removed_indices_tensor = reinterpret_cast(filter.addr); auto removed_indices = cuvs::core::from_dlpack(removed_indices_tensor); - cuvs::core::bitset removed_indices_bitset( - *res_ptr, removed_indices, index_ptr->dataset().extent(0)); - auto bitset_filter_obj = - cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset.view()); + cuvs::core::bitset_view removed_indices_bitset( + removed_indices, index_ptr->dataset().extent(0)); + auto bitset_filter_obj = cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset); cuvs::neighbors::cagra::search(*res_ptr, search_params, *index_ptr, @@ -138,7 +137,7 @@ void _search(cuvsResources_t res, distances_mds, bitset_filter_obj); } else { - RAFT_FAIL("Unsupported prefilter type: BITMAP"); + RAFT_FAIL("Unsupported filter type: BITMAP"); } } diff --git a/cpp/test/neighbors/ann_cagra_c.cu b/cpp/test/neighbors/ann_cagra_c.cu index 1a38d7b5ef..7315813cc9 100644 --- a/cpp/test/neighbors/ann_cagra_c.cu +++ b/cpp/test/neighbors/ann_cagra_c.cu @@ -34,7 +34,7 @@ float queries[4][2] = {{0.48216683, 0.0428398}, {0.51260436, 0.2643005}, {0.05198065, 0.5789965}}; -int64_t filter[2] = {1, 2}; +uint32_t filter[1] = {0b1001}; // index 1 and 2 are removed uint32_t neighbors_exp[4] = {3, 0, 3, 1}; float distances_exp[4] = {0.03878258, 0.12472608, 0.04776672, 0.15224178}; @@ -114,15 +114,15 @@ TEST(CagraC, BuildSearch) distances_tensor.dl_tensor.shape = distances_shape; distances_tensor.dl_tensor.strides = nullptr; - cuvsFilter prefilter; - prefilter.type = NO_FILTER; - prefilter.addr = (uintptr_t)NULL; + cuvsFilter filter; + filter.type = NO_FILTER; + filter.addr = (uintptr_t)NULL; // search index cuvsCagraSearchParams_t search_params; cuvsCagraSearchParamsCreate(&search_params); cuvsCagraSearch( - res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, prefilter); + res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter); // verify output ASSERT_TRUE( @@ -210,8 +210,8 @@ TEST(CagraC, BuildSearchFiltered) distances_tensor.dl_tensor.strides = nullptr; // create filter DLTensor - rmm::device_uvector filter_d(2, stream); - raft::copy(filter_d.data(), (int64_t*)filter, 2, stream); + rmm::device_uvector filter_d(1, stream); + raft::copy(filter_d.data(), filter, 1, stream); cuvsFilter filter; @@ -219,10 +219,10 @@ TEST(CagraC, BuildSearchFiltered) filter_tensor.dl_tensor.data = filter_d.data(); filter_tensor.dl_tensor.device.device_type = kDLCUDA; filter_tensor.dl_tensor.ndim = 1; - filter_tensor.dl_tensor.dtype.code = kDLInt; - filter_tensor.dl_tensor.dtype.bits = 64; + filter_tensor.dl_tensor.dtype.code = kDLUInt; + filter_tensor.dl_tensor.dtype.bits = 32; filter_tensor.dl_tensor.dtype.lanes = 1; - int64_t filter_shape[1] = {2}; + int64_t filter_shape[1] = {1}; filter_tensor.dl_tensor.shape = filter_shape; filter_tensor.dl_tensor.strides = nullptr; @@ -234,7 +234,6 @@ TEST(CagraC, BuildSearchFiltered) cuvsCagraSearchParamsCreate(&search_params); cuvsCagraSearch( res, search_params, index, &queries_tensor, &neighbors_tensor, &distances_tensor, filter); - // verify output ASSERT_TRUE(cuvs::devArrMatchHost( neighbors_exp_filtered, neighbors_d.data(), 4, cuvs::Compare())); From b0e8122945259a310c6c6185f8d9514dc255b728 Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Fri, 3 Jan 2025 22:08:16 +0000 Subject: [PATCH 08/16] python (WIP) --- python/cuvs/cuvs/neighbors/cagra/cagra.pxd | 5 +- python/cuvs/cuvs/neighbors/cagra/cagra.pyx | 11 +- .../cuvs/cuvs/neighbors/filters/__init__.py | 4 +- .../cuvs/cuvs/neighbors/filters/filters.pyx | 49 ++++++++ python/cuvs/cuvs/test/test_cagra.py | 108 +++++++++++++++++- 5 files changed, 170 insertions(+), 7 deletions(-) diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index bba5a91a8e..0efa140495 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -28,7 +28,7 @@ from libcpp cimport bool from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t from cuvs.common.cydlpack cimport DLDataType, DLManagedTensor - +from cuvs.neighbors.filters.filters cimport cuvsFilter cdef extern from "cuvs/neighbors/cagra.h" nogil: @@ -113,7 +113,8 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil: cuvsCagraIndex_t index, DLManagedTensor* queries, DLManagedTensor* neighbors, - DLManagedTensor* distances) except + + DLManagedTensor* distances, + cuvsFilter filter) except + cuvsError_t cuvsCagraSerialize(cuvsResources_t res, const char * filename, diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index 752aef7415..f7a542f7b3 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -47,6 +47,8 @@ from libc.stdint cimport ( from cuvs.common.exceptions import check_cuvs +from cuvs.neighbors.filters import no_filter + cdef class CompressionParams: """ @@ -480,7 +482,8 @@ def search(SearchParams search_params, k, neighbors=None, distances=None, - resources=None): + resources=None, + filter=None): """ Find the k nearest neighbors for each query. @@ -553,6 +556,9 @@ def search(SearchParams search_params, _check_input_array(distances_cai, [np.dtype('float32')], exp_rows=n_queries, exp_cols=k) + if filter is None: + filter = no_filter() + cdef cuvsCagraSearchParams* params = &search_params.params cdef cydlpack.DLManagedTensor* queries_dlpack = \ cydlpack.dlpack_c(queries_cai) @@ -569,7 +575,8 @@ def search(SearchParams search_params, index.index, queries_dlpack, neighbors_dlpack, - distances_dlpack + distances_dlpack, + filter.prefilter )) return (distances, neighbors) diff --git a/python/cuvs/cuvs/neighbors/filters/__init__.py b/python/cuvs/cuvs/neighbors/filters/__init__.py index 2ad118965a..0ddf809c96 100644 --- a/python/cuvs/cuvs/neighbors/filters/__init__.py +++ b/python/cuvs/cuvs/neighbors/filters/__init__.py @@ -13,6 +13,6 @@ # limitations under the License. -from .filters import Prefilter, from_bitmap, no_filter +from .filters import Prefilter, from_bitmap, from_bitset, no_filter -__all__ = ["no_filter", "from_bitmap", "Prefilter"] +__all__ = ["no_filter", "from_bitmap", "from_bitset", "Prefilter"] diff --git a/python/cuvs/cuvs/neighbors/filters/filters.pyx b/python/cuvs/cuvs/neighbors/filters/filters.pyx index 9bc2a905ce..c3ef258a60 100644 --- a/python/cuvs/cuvs/neighbors/filters/filters.pyx +++ b/python/cuvs/cuvs/neighbors/filters/filters.pyx @@ -95,3 +95,52 @@ def from_bitmap(bitmap): filter.addr = bitmap_dlpack return Prefilter(filter, parent=bitmap) + +def from_bitset(bitset): + """ + Create a pre-filter from an array with type of uint32. + + Parameters + ---------- + bitmap : numpy.ndarray + An array with type of `uint32` where each bit in the array corresponds + to if a sample and query pair is greenlit (not filtered) or filtered. + The array is row-major, meaning the bits are ordered by rows first. + Each bit in a `uint32` element represents a different sample-query + pair. + + - Bit value of 1: The sample-query pair is greenlit (allowed). + - Bit value of 0: The sample-query pair is filtered. + + Returns + ------- + filter : cuvs.neighbors.filters.Prefilter + An instance of `Prefilter` that can be used to filter neighbors + based on the given bitmap. + {resources_docstring} + + Examples + -------- + + >>> import cupy as cp + >>> import numpy as np + >>> from cuvs.neighbors import filters + >>> + >>> n_samples = 50000 + >>> n_queries = 1000 + >>> + >>> n_bitmap = np.ceil(n_samples * n_queries / 32).astype(int) + >>> bitmap = cp.random.randint(1, 100, size=(n_bitmap,), dtype=cp.uint32) + >>> prefilter = filters.from_bitmap(bitmap) + """ + bitset_cai = wrap_array(bitset) + _check_input_array(bitset_cai, [np.dtype('uint32')]) + + cdef cydlpack.DLManagedTensor* bitset_dlpack = \ + cydlpack.dlpack_c(bitset_cai) + + cdef cuvsFilter filter + filter.type = BITSET + filter.addr = bitset_dlpack + + return Prefilter(filter, parent=bitset) diff --git a/python/cuvs/cuvs/test/test_cagra.py b/python/cuvs/cuvs/test/test_cagra.py index 56e132c23c..6974d59cb4 100644 --- a/python/cuvs/cuvs/test/test_cagra.py +++ b/python/cuvs/cuvs/test/test_cagra.py @@ -18,8 +18,9 @@ from pylibraft.common import device_ndarray from sklearn.neighbors import NearestNeighbors from sklearn.preprocessing import normalize +from scipy.spatial.distance import cdist -from cuvs.neighbors import cagra +from cuvs.neighbors import cagra, filters from cuvs.test.ann_utils import calc_recall, generate_data @@ -166,6 +167,111 @@ def test_cagra_dataset_dtype_host_device( }, ], ) + + +def create_sparse_bitset(n_size, sparsity): + """Create a sparse bitset array for testing filtering""" + bits_per_uint32 = 32 + num_bits = n_size + num_uint32s = (num_bits + bits_per_uint32 - 1) // bits_per_uint32 + num_ones = int(num_bits * sparsity) + + array = np.zeros(num_uint32s, dtype=np.uint32) + indices = np.random.choice(num_bits, num_ones, replace=False) + + for index in indices: + i = index // bits_per_uint32 + bit_position = index % bits_per_uint32 + array[i] |= 1 << bit_position + + return array + + +@pytest.mark.parametrize("n_rows", [1000, 5000]) +@pytest.mark.parametrize("n_cols", [10, 50]) +@pytest.mark.parametrize("n_queries", [10, 100]) +@pytest.mark.parametrize("k", [10, 20]) +@pytest.mark.parametrize("sparsity", [0.2, 0.5]) +def test_filtered_cagra( + n_rows, + n_cols, + n_queries, + k, + sparsity, +): + """Test CAGRA index with filtering using bitset""" + # Generate test data + dataset = generate_data((n_rows, n_cols), np.float32) + queries = generate_data((n_queries, n_cols), np.float32) + + # Create bitset for filtering + bitset = create_sparse_bitset(n_rows, sparsity) + + # Convert dataset and queries to device arrays + dataset_device = device_ndarray(dataset) + queries_device = device_ndarray(queries) + bitset_device = device_ndarray(bitset) + + # Build index + build_params = cagra.IndexParams( + metric="euclidean", + intermediate_graph_degree=64, + graph_degree=32, + build_algo="nn_descent", + ) + index = cagra.build(build_params, dataset_device) + + # Create filter + prefilter = filters.from_bitset(bitset_device) + + # Search with filter + out_idx = np.zeros((n_queries, k), dtype=np.uint32) + out_dist = np.zeros((n_queries, k), dtype=np.float32) + out_idx_device = device_ndarray(out_idx) + out_dist_device = device_ndarray(out_dist) + + search_params = cagra.SearchParams() + ret_distances, ret_indices = cagra.search( + search_params, + index, + queries_device, + k, + neighbors=out_idx_device, + distances=out_dist_device, + # filter=prefilter, + ) + + # Convert bitset to bool array for validation + bitset_as_uint8 = bitset.view(np.uint8) + bool_filter = np.unpackbits(bitset_as_uint8) + bool_filter = bool_filter.reshape(-1, 4, 8) + bool_filter = np.flip(bool_filter, axis=2) + bool_filter = bool_filter.reshape(-1)[:n_rows] + bool_filter = np.logical_not(bool_filter) # Flip so True means filtered + + # Get filtered dataset for reference calculation + non_filtered_mask = ~bool_filter + filtered_dataset = dataset[non_filtered_mask] + + # Calculate reference values with sklearn on filtered dataset + nn_skl = NearestNeighbors( + n_neighbors=k, algorithm="brute", metric="euclidean" + ) + nn_skl.fit(filtered_dataset) + skl_idx = nn_skl.kneighbors(queries, return_distance=False) + + # Get actual results + actual_indices = out_idx_device.copy_to_host() + actual_distances = out_dist_device.copy_to_host() + # Verify filtering - no filtered indices should be in results + filtered_indices = np.where(bool_filter)[0] + # for i in range(n_queries): + # assert not np.intersect1d(filtered_indices, actual_indices[i]).size + + # Verify recall compared to sklearn reference + recall = calc_recall(actual_indices, skl_idx) + assert recall > 0.7 + def test_cagra_index_params(params): # Note that inner_product tests use normalized input which we cannot # represent in int8, therefore we test only sqeuclidean metric here. From afbcf741d5f1a91c31eae0e1bd74a638dc159556 Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Sat, 4 Jan 2025 11:05:12 +0000 Subject: [PATCH 09/16] fix python test --- python/cuvs/cuvs/test/test_cagra.py | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/python/cuvs/cuvs/test/test_cagra.py b/python/cuvs/cuvs/test/test_cagra.py index 6974d59cb4..4d7afc9239 100644 --- a/python/cuvs/cuvs/test/test_cagra.py +++ b/python/cuvs/cuvs/test/test_cagra.py @@ -170,7 +170,6 @@ def test_cagra_dataset_dtype_host_device( def create_sparse_bitset(n_size, sparsity): - """Create a sparse bitset array for testing filtering""" bits_per_uint32 = 32 num_bits = n_size num_uint32s = (num_bits + bits_per_uint32 - 1) // bits_per_uint32 @@ -199,20 +198,15 @@ def test_filtered_cagra( k, sparsity, ): - """Test CAGRA index with filtering using bitset""" - # Generate test data dataset = generate_data((n_rows, n_cols), np.float32) queries = generate_data((n_queries, n_cols), np.float32) - # Create bitset for filtering bitset = create_sparse_bitset(n_rows, sparsity) - # Convert dataset and queries to device arrays dataset_device = device_ndarray(dataset) queries_device = device_ndarray(queries) bitset_device = device_ndarray(bitset) - # Build index build_params = cagra.IndexParams( metric="euclidean", intermediate_graph_degree=64, @@ -221,10 +215,8 @@ def test_filtered_cagra( ) index = cagra.build(build_params, dataset_device) - # Create filter prefilter = filters.from_bitset(bitset_device) - # Search with filter out_idx = np.zeros((n_queries, k), dtype=np.uint32) out_dist = np.zeros((n_queries, k), dtype=np.float32) out_idx_device = device_ndarray(out_idx) @@ -238,7 +230,7 @@ def test_filtered_cagra( k, neighbors=out_idx_device, distances=out_dist_device, - # filter=prefilter, + filter=prefilter, ) # Convert bitset to bool array for validation @@ -263,13 +255,22 @@ def test_filtered_cagra( # Get actual results actual_indices = out_idx_device.copy_to_host() actual_distances = out_dist_device.copy_to_host() + + filtered_idx_map = np.cumsum(~bool_filter) - 1 # -1 because cumsum starts at 1 + + # Map CAGRA indices to filtered space + mapped_actual_indices = np.take(filtered_idx_map, + actual_indices, + mode='clip') + # Verify filtering - no filtered indices should be in results filtered_indices = np.where(bool_filter)[0] - # for i in range(n_queries): - # assert not np.intersect1d(filtered_indices, actual_indices[i]).size + for i in range(n_queries): + assert not np.intersect1d(filtered_indices, actual_indices[i]).size + + # Now compare with sklearn results + recall = calc_recall(mapped_actual_indices, skl_idx) - # Verify recall compared to sklearn reference - recall = calc_recall(actual_indices, skl_idx) assert recall > 0.7 def test_cagra_index_params(params): From f6b80684a525e1becba63bd9311874c53a322f2a Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Sat, 4 Jan 2025 11:24:01 +0000 Subject: [PATCH 10/16] rust + c example --- examples/c/src/cagra_c_example.c | 12 ++++++++---- rust/cuvs/src/cagra/index.rs | 9 ++++++++- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/examples/c/src/cagra_c_example.c b/examples/c/src/cagra_c_example.c index fdcbbf5712..83957aabca 100644 --- a/examples/c/src/cagra_c_example.c +++ b/examples/c/src/cagra_c_example.c @@ -67,9 +67,9 @@ void cagra_build_search_simple() { // Allocate memory for `queries`, `neighbors` and `distances` output uint32_t *neighbors; float *distances, *queries_d; - cuvsRMMAlloc(res, (void**) &queries_d, sizeof(float) * n_queries * n_cols); - cuvsRMMAlloc(res, (void**) &neighbors, sizeof(uint32_t) * n_queries * topk); - cuvsRMMAlloc(res, (void**) &distances, sizeof(float) * n_queries * topk); + cuvsRMMAlloc(res, (void **)&queries_d, sizeof(float) * n_queries * n_cols); + cuvsRMMAlloc(res, (void **)&neighbors, sizeof(uint32_t) * n_queries * topk); + cuvsRMMAlloc(res, (void **)&distances, sizeof(float) * n_queries * topk); // Use DLPack to represent `queries`, `neighbors` and `distances` as tensors cudaMemcpy(queries_d, queries, sizeof(float) * 4 * 2, cudaMemcpyDefault); @@ -111,8 +111,12 @@ void cagra_build_search_simple() { cuvsCagraSearchParams_t search_params; cuvsCagraSearchParamsCreate(&search_params); + cuvsFilter filter; + filter.type = NO_FILTER; + filter.addr = (uintptr_t)NULL; + cuvsCagraSearch(res, search_params, index, &queries_tensor, &neighbors_tensor, - &distances_tensor); + &distances_tensor, filter); // print results uint32_t *neighbors_h = diff --git a/rust/cuvs/src/cagra/index.rs b/rust/cuvs/src/cagra/index.rs index 959959f608..bf316b4d7a 100644 --- a/rust/cuvs/src/cagra/index.rs +++ b/rust/cuvs/src/cagra/index.rs @@ -78,6 +78,11 @@ impl Index { distances: &ManagedTensor, ) -> Result<()> { unsafe { + let prefilter = ffi::cuvsFilter { + addr: 0, + type_: ffi::cuvsFilterType::NO_FILTER, + }; + check_cuvs(ffi::cuvsCagraSearch( res.0, params.0, @@ -85,6 +90,7 @@ impl Index { queries.as_ptr(), neighbors.as_ptr(), distances.as_ptr(), + prefilter, )) } } @@ -167,7 +173,8 @@ mod tests { #[test] fn test_cagra_compression() { use crate::cagra::CompressionParams; - let build_params = IndexParams::new().unwrap() + let build_params = IndexParams::new() + .unwrap() .set_compression(CompressionParams::new().unwrap()); test_cagra(build_params); } From 9a64945fb28183e74b99b340d98575be4bb6f56f Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Wed, 8 Jan 2025 16:23:40 +0000 Subject: [PATCH 11/16] documentation, formatting --- python/cuvs/cuvs/neighbors/cagra/cagra.pxd | 2 +- python/cuvs/cuvs/neighbors/cagra/cagra.pyx | 4 +++- .../cuvs/cuvs/neighbors/filters/filters.pyx | 22 ++++++++---------- python/cuvs/cuvs/test/test_cagra.py | 23 ++++++++++--------- 4 files changed, 26 insertions(+), 25 deletions(-) diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd index 326602bc01..fba7e3d1ee 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd @@ -28,8 +28,8 @@ from libcpp cimport bool from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t from cuvs.common.cydlpack cimport DLDataType, DLManagedTensor -from cuvs.neighbors.filters.filters cimport cuvsFilter from cuvs.distance_type cimport cuvsDistanceType +from cuvs.neighbors.filters.filters cimport cuvsFilter cdef extern from "cuvs/neighbors/cagra.h" nogil: diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx index 95ebc1895f..f62563f610 100644 --- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx +++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx @@ -48,7 +48,6 @@ from libc.stdint cimport ( ) from cuvs.common.exceptions import check_cuvs - from cuvs.neighbors.filters import no_filter @@ -506,6 +505,9 @@ def search(SearchParams search_params, distances : Optional CUDA array interface compliant matrix shape (n_queries, k) If supplied, the distances to the neighbors will be written here in-place. (default None) + filter: Optional cuvs.neighbors.cuvsFilter can be used to filter + neighbors based on a given bitset. + (default None) {resources_docstring} Examples diff --git a/python/cuvs/cuvs/neighbors/filters/filters.pyx b/python/cuvs/cuvs/neighbors/filters/filters.pyx index c3ef258a60..7bc6f9dae2 100644 --- a/python/cuvs/cuvs/neighbors/filters/filters.pyx +++ b/python/cuvs/cuvs/neighbors/filters/filters.pyx @@ -22,7 +22,7 @@ from libc.stdint cimport uintptr_t from cuvs.common cimport cydlpack from cuvs.neighbors.common import _check_input_array -from .filters cimport BITMAP, NO_FILTER, cuvsFilter +from .filters cimport BITMAP, BITSET, NO_FILTER, cuvsFilter from pylibraft.common.cai_wrapper import wrap_array @@ -102,21 +102,19 @@ def from_bitset(bitset): Parameters ---------- - bitmap : numpy.ndarray + bitset : numpy.ndarray An array with type of `uint32` where each bit in the array corresponds - to if a sample and query pair is greenlit (not filtered) or filtered. - The array is row-major, meaning the bits are ordered by rows first. - Each bit in a `uint32` element represents a different sample-query - pair. + to if a sample is greenlit (not filtered) or filtered. + Each bit in a `uint32` element represents a different sample of the dataset. - - Bit value of 1: The sample-query pair is greenlit (allowed). - - Bit value of 0: The sample-query pair is filtered. + - Bit value of 1: The sample is greenlit (allowed). + - Bit value of 0: The sample pair is filtered. Returns ------- filter : cuvs.neighbors.filters.Prefilter An instance of `Prefilter` that can be used to filter neighbors - based on the given bitmap. + based on the given bitset. {resources_docstring} Examples @@ -129,9 +127,9 @@ def from_bitset(bitset): >>> n_samples = 50000 >>> n_queries = 1000 >>> - >>> n_bitmap = np.ceil(n_samples * n_queries / 32).astype(int) - >>> bitmap = cp.random.randint(1, 100, size=(n_bitmap,), dtype=cp.uint32) - >>> prefilter = filters.from_bitmap(bitmap) + >>> n_bitset = np.ceil(n_samples / 32).astype(int) + >>> bitset = cp.random.randint(1, 100, size=(n_bitset,), dtype=cp.uint32) + >>> prefilter = filters.from_bitset(bitset) """ bitset_cai = wrap_array(bitset) _check_input_array(bitset_cai, [np.dtype('uint32')]) diff --git a/python/cuvs/cuvs/test/test_cagra.py b/python/cuvs/cuvs/test/test_cagra.py index 608796fd44..650eb7b5d1 100644 --- a/python/cuvs/cuvs/test/test_cagra.py +++ b/python/cuvs/cuvs/test/test_cagra.py @@ -16,9 +16,9 @@ import numpy as np import pytest from pylibraft.common import device_ndarray +from scipy.spatial.distance import cdist from sklearn.neighbors import NearestNeighbors from sklearn.preprocessing import normalize -from scipy.spatial.distance import cdist from cuvs.neighbors import cagra, filters from cuvs.test.ann_utils import calc_recall, generate_data @@ -169,11 +169,9 @@ def test_cagra_dataset_dtype_host_device( }, ], ) - - def create_sparse_bitset(n_size, sparsity): bits_per_uint32 = 32 - num_bits = n_size + num_bits = n_size num_uint32s = (num_bits + bits_per_uint32 - 1) // bits_per_uint32 num_ones = int(num_bits * sparsity) @@ -257,14 +255,16 @@ def test_filtered_cagra( # Get actual results actual_indices = out_idx_device.copy_to_host() actual_distances = out_dist_device.copy_to_host() - - filtered_idx_map = np.cumsum(~bool_filter) - 1 # -1 because cumsum starts at 1 + + filtered_idx_map = ( + np.cumsum(~bool_filter) - 1 + ) # -1 because cumsum starts at 1 # Map CAGRA indices to filtered space - mapped_actual_indices = np.take(filtered_idx_map, - actual_indices, - mode='clip') - + mapped_actual_indices = np.take( + filtered_idx_map, actual_indices, mode="clip" + ) + # Verify filtering - no filtered indices should be in results filtered_indices = np.where(bool_filter)[0] for i in range(n_queries): @@ -274,7 +274,8 @@ def test_filtered_cagra( recall = calc_recall(mapped_actual_indices, skl_idx) assert recall > 0.7 - + + def test_cagra_index_params(params): # Note that inner_product tests use normalized input which we cannot # represent in int8, therefore we test only sqeuclidean metric here. From 8ac5b7fb85749939acc927d53d44754bf0ea23e8 Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Fri, 10 Jan 2025 16:11:34 +0000 Subject: [PATCH 12/16] fix style issues --- cpp/include/cuvs/neighbors/cagra.h | 2 +- python/cuvs/cuvs/neighbors/filters/filters.pyx | 8 +++++--- python/cuvs/cuvs/test/test_cagra.py | 2 -- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index 81a7f02150..f519baaba3 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -17,8 +17,8 @@ #pragma once #include -#include #include +#include #include #include #include diff --git a/python/cuvs/cuvs/neighbors/filters/filters.pyx b/python/cuvs/cuvs/neighbors/filters/filters.pyx index 7bc6f9dae2..0de1c51615 100644 --- a/python/cuvs/cuvs/neighbors/filters/filters.pyx +++ b/python/cuvs/cuvs/neighbors/filters/filters.pyx @@ -96,6 +96,7 @@ def from_bitmap(bitmap): return Prefilter(filter, parent=bitmap) + def from_bitset(bitset): """ Create a pre-filter from an array with type of uint32. @@ -103,9 +104,10 @@ def from_bitset(bitset): Parameters ---------- bitset : numpy.ndarray - An array with type of `uint32` where each bit in the array corresponds - to if a sample is greenlit (not filtered) or filtered. - Each bit in a `uint32` element represents a different sample of the dataset. + An array with type of `uint32` where each bit in the array + corresponds to if a sample is greenlit (not filtered) or filtered. + Each bit in a `uint32` element represents a different sample of + the dataset. - Bit value of 1: The sample is greenlit (allowed). - Bit value of 0: The sample pair is filtered. diff --git a/python/cuvs/cuvs/test/test_cagra.py b/python/cuvs/cuvs/test/test_cagra.py index 650eb7b5d1..f6ad8001f8 100644 --- a/python/cuvs/cuvs/test/test_cagra.py +++ b/python/cuvs/cuvs/test/test_cagra.py @@ -16,7 +16,6 @@ import numpy as np import pytest from pylibraft.common import device_ndarray -from scipy.spatial.distance import cdist from sklearn.neighbors import NearestNeighbors from sklearn.preprocessing import normalize @@ -254,7 +253,6 @@ def test_filtered_cagra( # Get actual results actual_indices = out_idx_device.copy_to_host() - actual_distances = out_dist_device.copy_to_host() filtered_idx_map = ( np.cumsum(~bool_filter) - 1 From 9d8ee2e6d0b8f287801eda2201f219ab99b0144d Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Tue, 14 Jan 2025 20:17:07 +0000 Subject: [PATCH 13/16] fix whitespace issue --- python/cuvs/cuvs/neighbors/filters/filters.pyx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cuvs/cuvs/neighbors/filters/filters.pyx b/python/cuvs/cuvs/neighbors/filters/filters.pyx index 0de1c51615..16042e9662 100644 --- a/python/cuvs/cuvs/neighbors/filters/filters.pyx +++ b/python/cuvs/cuvs/neighbors/filters/filters.pyx @@ -104,9 +104,9 @@ def from_bitset(bitset): Parameters ---------- bitset : numpy.ndarray - An array with type of `uint32` where each bit in the array + An array with type of `uint32` where each bit in the array corresponds to if a sample is greenlit (not filtered) or filtered. - Each bit in a `uint32` element represents a different sample of + Each bit in a `uint32` element represents a different sample of the dataset. - Bit value of 1: The sample is greenlit (allowed). From 7e2b80322ffc0c03c1595de6a3e770486c47cce2 Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Thu, 16 Jan 2025 23:09:36 +0000 Subject: [PATCH 14/16] fix errors (wip) --- cpp/include/cuvs/neighbors/cagra.h | 2 +- python/cuvs/cuvs/test/test_cagra.py | 72 ++++++++++++++--------------- 2 files changed, 37 insertions(+), 37 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.h b/cpp/include/cuvs/neighbors/cagra.h index f519baaba3..3e17f1d0fe 100644 --- a/cpp/include/cuvs/neighbors/cagra.h +++ b/cpp/include/cuvs/neighbors/cagra.h @@ -389,7 +389,7 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res, * @param[in] queries DLManagedTensor* queries dataset to search * @param[out] neighbors DLManagedTensor* output `k` neighbors for queries * @param[out] distances DLManagedTensor* output `k` distances for queries - * @param[in] prefilter cuvsFilter input prefilter that can be used + * @param[in] filter cuvsFilter input filter that can be used to filter queries and neighbors based on the given bitset. */ cuvsError_t cuvsCagraSearch(cuvsResources_t res, diff --git a/python/cuvs/cuvs/test/test_cagra.py b/python/cuvs/cuvs/test/test_cagra.py index f6ad8001f8..4f6a1538dc 100644 --- a/python/cuvs/cuvs/test/test_cagra.py +++ b/python/cuvs/cuvs/test/test_cagra.py @@ -139,35 +139,6 @@ def test_cagra_dataset_dtype_host_device( ) -@pytest.mark.parametrize( - "params", - [ - { - "intermediate_graph_degree": 64, - "graph_degree": 32, - "add_data_on_build": True, - "k": 1, - "metric": "sqeuclidean", - "build_algo": "ivf_pq", - }, - { - "intermediate_graph_degree": 32, - "graph_degree": 16, - "add_data_on_build": False, - "k": 5, - "metric": "sqeuclidean", - "build_algo": "ivf_pq", - }, - { - "intermediate_graph_degree": 128, - "graph_degree": 32, - "add_data_on_build": True, - "k": 10, - "metric": "inner_product", - "build_algo": "nn_descent", - }, - ], -) def create_sparse_bitset(n_size, sparsity): bits_per_uint32 = 32 num_bits = n_size @@ -185,10 +156,10 @@ def create_sparse_bitset(n_size, sparsity): return array -@pytest.mark.parametrize("n_rows", [1000, 5000]) -@pytest.mark.parametrize("n_cols", [10, 50]) -@pytest.mark.parametrize("n_queries", [10, 100]) -@pytest.mark.parametrize("k", [10, 20]) +@pytest.mark.parametrize("n_rows", [10000]) +@pytest.mark.parametrize("n_cols", [10]) +@pytest.mark.parametrize("n_queries", [10]) +@pytest.mark.parametrize("k", [10]) @pytest.mark.parametrize("sparsity", [0.2, 0.5]) def test_filtered_cagra( n_rows, @@ -207,14 +178,14 @@ def test_filtered_cagra( bitset_device = device_ndarray(bitset) build_params = cagra.IndexParams( - metric="euclidean", + metric="sqeuclidean", intermediate_graph_degree=64, graph_degree=32, build_algo="nn_descent", ) index = cagra.build(build_params, dataset_device) - prefilter = filters.from_bitset(bitset_device) + filter_ = filters.from_bitset(bitset_device) out_idx = np.zeros((n_queries, k), dtype=np.uint32) out_dist = np.zeros((n_queries, k), dtype=np.float32) @@ -229,7 +200,7 @@ def test_filtered_cagra( k, neighbors=out_idx_device, distances=out_dist_device, - filter=prefilter, + filter=filter_, ) # Convert bitset to bool array for validation @@ -274,6 +245,35 @@ def test_filtered_cagra( assert recall > 0.7 +@pytest.mark.parametrize( + "params", + [ + { + "intermediate_graph_degree": 64, + "graph_degree": 32, + "add_data_on_build": True, + "k": 1, + "metric": "sqeuclidean", + "build_algo": "ivf_pq", + }, + { + "intermediate_graph_degree": 32, + "graph_degree": 16, + "add_data_on_build": False, + "k": 5, + "metric": "sqeuclidean", + "build_algo": "ivf_pq", + }, + { + "intermediate_graph_degree": 128, + "graph_degree": 32, + "add_data_on_build": True, + "k": 10, + "metric": "inner_product", + "build_algo": "nn_descent", + }, + ], +) def test_cagra_index_params(params): # Note that inner_product tests use normalized input which we cannot # represent in int8, therefore we test only sqeuclidean metric here. From a307c5750cf86441e0ac0a279621fb4ce0c4a498 Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Fri, 17 Jan 2025 10:35:16 +0000 Subject: [PATCH 15/16] cleanup --- python/cuvs/cuvs/test/test_cagra.py | 25 ++++++------------------- 1 file changed, 6 insertions(+), 19 deletions(-) diff --git a/python/cuvs/cuvs/test/test_cagra.py b/python/cuvs/cuvs/test/test_cagra.py index 4f6a1538dc..f2bb3a39f0 100644 --- a/python/cuvs/cuvs/test/test_cagra.py +++ b/python/cuvs/cuvs/test/test_cagra.py @@ -156,17 +156,13 @@ def create_sparse_bitset(n_size, sparsity): return array -@pytest.mark.parametrize("n_rows", [10000]) -@pytest.mark.parametrize("n_cols", [10]) -@pytest.mark.parametrize("n_queries", [10]) -@pytest.mark.parametrize("k", [10]) -@pytest.mark.parametrize("sparsity", [0.2, 0.5]) +@pytest.mark.parametrize("sparsity", [0.2, 0.5, 0.7, 1.0]) def test_filtered_cagra( - n_rows, - n_cols, - n_queries, - k, sparsity, + n_rows = 10000, + n_cols = 10, + n_queries = 10, + k = 10, ): dataset = generate_data((n_rows, n_cols), np.float32) queries = generate_data((n_queries, n_cols), np.float32) @@ -177,12 +173,7 @@ def test_filtered_cagra( queries_device = device_ndarray(queries) bitset_device = device_ndarray(bitset) - build_params = cagra.IndexParams( - metric="sqeuclidean", - intermediate_graph_degree=64, - graph_degree=32, - build_algo="nn_descent", - ) + build_params = cagra.IndexParams() index = cagra.build(build_params, dataset_device) filter_ = filters.from_bitset(bitset_device) @@ -215,14 +206,12 @@ def test_filtered_cagra( non_filtered_mask = ~bool_filter filtered_dataset = dataset[non_filtered_mask] - # Calculate reference values with sklearn on filtered dataset nn_skl = NearestNeighbors( n_neighbors=k, algorithm="brute", metric="euclidean" ) nn_skl.fit(filtered_dataset) skl_idx = nn_skl.kneighbors(queries, return_distance=False) - # Get actual results actual_indices = out_idx_device.copy_to_host() filtered_idx_map = ( @@ -234,12 +223,10 @@ def test_filtered_cagra( filtered_idx_map, actual_indices, mode="clip" ) - # Verify filtering - no filtered indices should be in results filtered_indices = np.where(bool_filter)[0] for i in range(n_queries): assert not np.intersect1d(filtered_indices, actual_indices[i]).size - # Now compare with sklearn results recall = calc_recall(mapped_actual_indices, skl_idx) assert recall > 0.7 From 833165619ab63347a1b1788513dfc4003b666a68 Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Thu, 23 Jan 2025 21:58:34 +0000 Subject: [PATCH 16/16] fix style check --- python/cuvs/cuvs/test/test_cagra.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/cuvs/cuvs/test/test_cagra.py b/python/cuvs/cuvs/test/test_cagra.py index f2bb3a39f0..9efd489233 100644 --- a/python/cuvs/cuvs/test/test_cagra.py +++ b/python/cuvs/cuvs/test/test_cagra.py @@ -159,10 +159,10 @@ def create_sparse_bitset(n_size, sparsity): @pytest.mark.parametrize("sparsity", [0.2, 0.5, 0.7, 1.0]) def test_filtered_cagra( sparsity, - n_rows = 10000, - n_cols = 10, - n_queries = 10, - k = 10, + n_rows=10000, + n_cols=10, + n_queries=10, + k=10, ): dataset = generate_data((n_rows, n_cols), np.float32) queries = generate_data((n_queries, n_cols), np.float32)