From 59bb8ca9afeaddc0b7dce234f2d49df000254394 Mon Sep 17 00:00:00 2001 From: Ajit Mistry Date: Tue, 5 Nov 2024 13:45:25 +0000 Subject: [PATCH 1/7] . --- 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 14331ebbc..b629e5c1f 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 6985ff094..bb178de7c 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 599d2d842..9330cf573 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 2/7] 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 9330cf573..6e5008d58 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 3/7] 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 bb178de7c..c3631d40a 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 6e5008d58..2b6e7e5e2 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 4/7] 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 b629e5c1f..6abe43e38 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 25408ff01..ade29fa06 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 2b6e7e5e2..1a38d7b5e 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 5/7] . --- 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 6abe43e38..95e4ea260 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 6/7] 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 ade29fa06..dc9abe070 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 7/7] 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 dc9abe070..639f32d8e 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 1a38d7b5e..7315813cc 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()));