diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 66f139c090..84363a38b4 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -199,6 +199,7 @@ if(NOT BUILD_CPU_ONLY) include(cmake/thirdparty/get_raft.cmake) include(cmake/thirdparty/get_kvikio.cmake) include(cmake/thirdparty/get_cutlass.cmake) + rapids_cpm_package_override("${CMAKE_CURRENT_SOURCE_DIR}/cmake/patches/cuco_override.json") include(${rapids-cmake-dir}/cpm/cuco.cmake) rapids_cpm_cuco() include(cmake/thirdparty/get_rtcx.cmake) @@ -1394,6 +1395,8 @@ if(NOT BUILD_CPU_ONLY) src/cluster/spectral.cu src/core/bitset.cu src/core/bloom_filter.cu + src/core/roaring_allowlist.cu + src/neighbors/roaring_filter.cu src/core/omp_wrapper.cpp src/util/file_io.cpp src/util/host_memory.cpp diff --git a/cpp/cmake/patches/cuco_override.json b/cpp/cmake/patches/cuco_override.json new file mode 100644 index 0000000000..004b881cb2 --- /dev/null +++ b/cpp/cmake/patches/cuco_override.json @@ -0,0 +1,9 @@ +{ + "packages" : { + "cuco" : { + "version": "0.0.1", + "git_url": "https://github.com/NVIDIA/cuCollections.git", + "git_tag": "2027216da1ba687374222f54e0e196dd1ccaedc7" + } + } +} diff --git a/cpp/include/cuvs/core/roaring_allowlist.hpp b/cpp/include/cuvs/core/roaring_allowlist.hpp new file mode 100644 index 0000000000..af27678252 --- /dev/null +++ b/cpp/include/cuvs/core/roaring_allowlist.hpp @@ -0,0 +1,138 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include +#include + +#include +#include +#include + +namespace CUVS_EXPORT cuvs { +namespace core { + +/** + * @brief Non-owning device view of one immutable Roaring allowlist. + * + * The view contains an opaque pointer to an already initialized device-side cuCollections + * reference plus immutable shape and cardinality metadata. Creating or copying it is O(1) and + * performs no allocation, parsing, kernel launch, or synchronization. The owning + * @ref roaring_allowlist must outlive the view and every operation that uses it. + */ +class CUVS_EXPORT roaring_allowlist_view { + public: + roaring_allowlist_view() = default; + + [[nodiscard]] std::size_t dataset_rows() const noexcept { return dataset_rows_; } + [[nodiscard]] std::size_t cardinality() const noexcept { return cardinality_; } + [[nodiscard]] bool empty() const noexcept { return cardinality_ == 0; } + [[nodiscard]] bool valid() const noexcept { return valid_; } + + /** @brief Opaque device pointer to the initialized cuCollections reference, or null if empty. */ + [[nodiscard]] void const* device_reference() const noexcept { return device_reference_; } + + private: + friend class roaring_allowlist; + + roaring_allowlist_view(void const* device_reference, + std::size_t dataset_rows, + std::size_t cardinality) noexcept + : device_reference_(device_reference), + dataset_rows_(dataset_rows), + cardinality_(cardinality), + valid_(true) + { + } + + void const* device_reference_{}; + std::size_t dataset_rows_{}; + std::size_t cardinality_{}; + bool valid_{}; +}; + +/** + * @brief Owning immutable exact Roaring allowlist over CAGRA dataset-row IDs. + * + * Build an allowlist through @ref from_ids, then pass its zero-copy @ref view to a + * cuvs::neighbors::filtering::roaring_filter. A filter maps one such view to each query; owners + * remain independent and can therefore be reused across filters and queries. + * + * Construction sorts IDs on the GPU unless @p pre_sorted is true. Setting @p pre_sorted promises + * that IDs are already in strictly increasing order; this promise is not verified. IDs must be + * unique and smaller than @p dataset_rows. The encoded bytes and the initialized + * cuco::experimental::roaring_bitmap_ref are retained on the device. Creating a view + * never copies or reparses them, and CAGRA search performs no Roaring initialization. + * + * ID-based construction emits the standard portable 32-bit Roaring array and bitmap container + * forms. Each ID is partitioned by its high 16 bits; the low 16 bits are stored as an array for at + * most 4,096 values in a partition and as an 8 KiB bitmap otherwise. + * + * @see https://github.com/RoaringBitmap/RoaringFormatSpec + * @see https://github.com/NVIDIA/cuCollections/pull/839 + */ +class CUVS_EXPORT roaring_allowlist { + private: + struct impl; + + public: + using key_type = std::uint32_t; + + /** + * @brief Build one allowlist from host IDs. + * + * Host IDs are copied to the construction stream and then use the same device builder as the + * device overload. IDs must be smaller than dataset_rows; this precondition is not checked. + * Empty input is valid and rejects every candidate. The returned object is ready for same-stream + * use; cross-stream use requires an explicit dependency on the construction stream. + */ + static roaring_allowlist from_ids(raft::resources const& res, + std::size_t dataset_rows, + raft::host_vector_view ids, + bool pre_sorted = false); + + /** + * @brief Build one allowlist from device IDs. + * + * The input must remain valid until the construction stream reaches the enqueued work. IDs must + * be smaller than dataset_rows; this precondition is not checked. + * Temporary memory is O(cardinality + container count); no dataset-sized dense bitmap is used. + * The returned object is ready for same-stream use; cross-stream use requires an explicit + * dependency on the construction stream. + */ + static roaring_allowlist from_ids(raft::resources const& res, + std::size_t dataset_rows, + raft::device_vector_view ids, + bool pre_sorted = false); + + ~roaring_allowlist(); + + roaring_allowlist(roaring_allowlist const&) = delete; + roaring_allowlist& operator=(roaring_allowlist const&) = delete; + roaring_allowlist(roaring_allowlist&&) noexcept; + roaring_allowlist& operator=(roaring_allowlist&&) noexcept; + + [[nodiscard]] std::size_t dataset_rows() const noexcept; + [[nodiscard]] std::size_t cardinality() const noexcept; + [[nodiscard]] bool empty() const noexcept; + + /** @brief Total device bytes retained by the encoded allowlist and initialized reference. */ + [[nodiscard]] std::size_t size_bytes() const noexcept; + + /** @brief Return a zero-copy view. */ + [[nodiscard]] roaring_allowlist_view view() const noexcept; + + private: + explicit roaring_allowlist(std::unique_ptr impl) noexcept; + + std::unique_ptr impl_; +}; + +} // namespace core +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp index c1a73687c2..56180e3434 100644 --- a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp @@ -15,6 +15,7 @@ struct tag_u8 {}; struct tag_filter_none {}; struct tag_filter_bitset {}; struct tag_filter_bloom_filter {}; +struct tag_filter_roaring {}; struct tag_filter_udf {}; struct tag_bitset_u32 {}; diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 4ad9aea688..e9d3a6a43a 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -33,6 +33,7 @@ #include #include #include +#include #include #include #include @@ -43,7 +44,9 @@ namespace CUVS_EXPORT cuvs { namespace core { class bloom_filter; -} +class roaring_allowlist_view; +class roaring_allowlist; +} // namespace core namespace neighbors { /** * @addtogroup cagra_cpp_index_params @@ -1346,7 +1349,7 @@ namespace filtering { * @{ */ -enum class FilterType : int { None = 0, Bitmap = 1, Bitset = 2, Bloom = 3, UDF = 100 }; +enum class FilterType : int { None = 0, Bitmap = 1, Bitset = 2, Bloom = 3, Roaring = 4, UDF = 100 }; struct base_filter { ~base_filter() = default; @@ -1501,6 +1504,94 @@ struct bloom_filter : public base_filter { FilterType get_filter_type() const override { return FilterType::Bloom; } }; +/** + * @brief Reusable per-query mapping to immutable exact Roaring allowlists. + * + * Entry @c q selects view @c q. CAGRA retains candidate dataset row @c r when the selected + * allowlist contains @c r. Construction copies only already initialized device-reference pointers + * and empty flags into the filter payload; encoded bytes are neither copied nor parsed. Search + * therefore performs no Roaring allocation, initialization, synchronization, or preprocessing. + * + * @code{.cpp} + * auto first = cuvs::core::roaring_allowlist::from_ids( + * res, dataset_rows, + * raft::make_host_vector_view(first_ids.data(), + * first_ids.size())); + * auto second = cuvs::core::roaring_allowlist::from_ids( + * res, dataset_rows, + * raft::make_host_vector_view(second_ids.data(), + * second_ids.size())); + * std::array views{first.view(), second.view()}; + * auto filter = cuvs::neighbors::filtering::roaring_filter(res, views); + * @endcode + * + * Owners and views can be reused across filters and queries. This filter owns its mapping tables + * and device payload, but not the referenced owners, which must outlive the filter and all searches + * using it. Copies are cheap shared handles required by CAGRA query-offset wrappers. + * + * Roaring filters currently support direct @c cagra::search only. Dynamic batching can combine + * requests into a different query-row layout, and tiered search applies one filter to partitions + * with different row domains; both paths reject this filter type. + * + * @see cuvs::core::roaring_allowlist + * @see https://github.com/RoaringBitmap/RoaringFormatSpec + */ +struct roaring_filter : public base_filter { + private: + struct impl; + + public: + /** @brief Construct an invalid handle. It cannot be passed to CAGRA search. */ + roaring_filter() = default; + + /** + * @brief Materialize the query-to-allowlist device pointer table. + * + * @p allowlists must be nonempty, every view must be valid, and every view must have the same + * `dataset_rows()`. Query count is inferred from the span length. + */ + explicit roaring_filter(raft::resources const& res, + std::span allowlists); + + [[nodiscard]] bool valid() const noexcept; + [[nodiscard]] std::size_t num_queries() const noexcept; + [[nodiscard]] std::size_t dataset_rows() const noexcept; + [[nodiscard]] std::size_t cardinality(std::size_t query_id) const; + [[nodiscard]] bool empty(std::size_t query_id) const; + + /** + * @brief Conservative maximum rejected fraction among all query allowlists. + * + * CAGRA uses this precomputed value when `search_params::filtering_rate` is unset. Basing one + * batch-wide scalar on the sparsest query avoids under-provisioning that query, but a very sparse + * or empty allowlist can increase the search work performed for every query in the batch. Callers + * may set `search_params::filtering_rate` explicitly when another tradeoff is preferable. + */ + [[nodiscard]] float filtering_rate() const noexcept; + + /** @brief Device bytes owned by this mapping, excluding the referenced allowlists. */ + [[nodiscard]] std::size_t size_bytes() const noexcept; + + /** + * @brief Replace one query's allowlist pointer outside the search path. + * + * The replacement must have the same `dataset_rows()`. The method copies one pointer and one + * empty flag to the device and synchronizes @p res before returning. Do not call it concurrently + * with a search, and keep the replacement owner alive for all subsequent searches. + */ + void set_allowlist(raft::resources const& res, + std::size_t query_id, + cuvs::core::roaring_allowlist_view replacement); + + /** @brief Internal device payload already prepared for the linked CAGRA predicate. */ + [[nodiscard]] void* device_payload() const noexcept; + + FilterType get_filter_type() const override { return FilterType::Roaring; } + + private: + std::shared_ptr impl_; +}; + /** * @brief JIT-LTO user-defined filter predicate. * diff --git a/cpp/include/cuvs/neighbors/dynamic_batching.hpp b/cpp/include/cuvs/neighbors/dynamic_batching.hpp index 720bfee4d0..1207d31ae9 100644 --- a/cpp/include/cuvs/neighbors/dynamic_batching.hpp +++ b/cpp/include/cuvs/neighbors/dynamic_batching.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -173,6 +173,7 @@ struct index : cuvs::neighbors::index { * @param[in] sample_filter * filtering function, if any, must be the same for all requests in a batch * (the pointer must be alive for the lifetime of the dynamic batching index) + * Roaring filters are not supported because batching changes the query-row mapping. */ template index(const raft::resources& res, diff --git a/cpp/include/cuvs/neighbors/tiered_index.hpp b/cpp/include/cuvs/neighbors/tiered_index.hpp index 27d0114087..015e48c7e2 100644 --- a/cpp/include/cuvs/neighbors/tiered_index.hpp +++ b/cpp/include/cuvs/neighbors/tiered_index.hpp @@ -189,7 +189,8 @@ void compact(raft::resources const& res, * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries, * k] * @param[in] sample_filter an optional device filter function object that greenlights samples - * for a given query. (none_sample_filter for no filtering) + * for a given query. (none_sample_filter for no filtering). Roaring filters are not supported + * because tiered partitions use different dataset-row domains. */ void search(raft::resources const& res, const cagra::search_params& search_params, diff --git a/cpp/src/core/roaring_allowlist.cu b/cpp/src/core/roaring_allowlist.cu new file mode 100644 index 0000000000..16631b56a0 --- /dev/null +++ b/cpp/src/core/roaring_allowlist.cu @@ -0,0 +1,173 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include "nvtx.hpp" + +#include + +#include +#include +#include +#include + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace cuvs::core { +namespace { + +/** + * cuCollections Roaring ownership and encoding + * ============================================= + * + * cuCollections constructs a standard portable 32-bit Roaring stream directly on the GPU. A + * 32-bit ID is divided into a high-16-bit container key and a low-16-bit value. For each key, at + * most 4,096 values are encoded as a sorted uint16 array; larger containers use an 8 KiB bitmap. + * ID construction does not emit run containers. + * + * The owning cuco object retains the encoded bytes. Its lightweight + * cuco::experimental::roaring_bitmap_ref parses the portable header once and then stores + * container-location metadata plus pointers into those bytes. cuVS copies that reference into a + * stable device allocation during construction. Views and CAGRA filters copy only its pointer, so + * search never copies or reparses the payload. + * + * @see https://github.com/RoaringBitmap/RoaringFormatSpec + * @see https://github.com/NVIDIA/cuCollections/pull/839 + */ + +using ref_type = cuco::experimental::roaring_bitmap_ref; +using cuco_bitmap_allocator = rmm::mr::polymorphic_allocator; +using cuco_bitmap_type = cuco::experimental::roaring_bitmap; + +void validate_dataset_rows(std::size_t dataset_rows) +{ + constexpr std::uint64_t kKeyDomain = std::uint64_t{1} << 32; + RAFT_EXPECTS(dataset_rows > 0, "dataset_rows must be greater than zero."); + RAFT_EXPECTS(static_cast(dataset_rows) <= kKeyDomain, + "dataset_rows exceeds the uint32_t Roaring key domain."); +} + +struct device_build_result { + std::unique_ptr owner; + std::size_t cardinality{}; +}; + +device_build_result build_from_device_ids( + raft::resources const& res, + raft::device_vector_view ids, + bool pre_sorted) +{ + common::nvtx::range build_scope("roaring_allowlist::build_from_ids"); + auto const stream = raft::resource::get_cuda_stream(res); + auto const size = static_cast(ids.extent(0)); + if (size == 0) { return {}; } + + cuco_bitmap_allocator allocator{}; + cuda::stream_ref cuco_stream{stream.value()}; + auto bitmap = pre_sorted ? cuco_bitmap_type::from_sorted_unique_indices( + ids.data_handle(), ids.data_handle() + size, allocator, cuco_stream) + : cuco_bitmap_type::from_indices( + ids.data_handle(), ids.data_handle() + size, allocator, cuco_stream); + auto owner = std::make_unique(std::move(bitmap)); + auto const cardinality = static_cast(owner->size()); + return {std::move(owner), cardinality}; +} + +} // namespace + +struct roaring_allowlist::impl { + std::unique_ptr owner; + std::optional host_reference; + rmm::device_uvector device_reference; + std::size_t cardinality_{}; + std::size_t dataset_rows_{}; + + impl(raft::resources const& res, std::size_t dataset_rows, device_build_result&& built) + : owner(std::move(built.owner)), + host_reference(owner ? std::make_optional(owner->ref()) : std::nullopt), + device_reference(owner ? 1 : 0, raft::resource::get_cuda_stream(res)), + cardinality_(built.cardinality), + dataset_rows_(dataset_rows) + { + static_assert(std::is_trivially_copyable_v); + if (host_reference) { + raft::copy( + device_reference.data(), &host_reference.value(), 1, raft::resource::get_cuda_stream(res)); + } + } + + [[nodiscard]] ref_type const* reference() const noexcept + { + return device_reference.size() == 0 ? nullptr : device_reference.data(); + } +}; + +roaring_allowlist::roaring_allowlist(std::unique_ptr impl) noexcept : impl_(std::move(impl)) +{ +} + +roaring_allowlist roaring_allowlist::from_ids( + raft::resources const& res, + std::size_t dataset_rows, + raft::host_vector_view ids, + bool pre_sorted) +{ + validate_dataset_rows(dataset_rows); + auto const size = static_cast(ids.extent(0)); + auto const stream = raft::resource::get_cuda_stream(res); + rmm::device_uvector device_ids(size, stream); + if (size != 0) { raft::copy(device_ids.data(), ids.data_handle(), size, stream); } + auto device_ids_view = raft::make_device_vector_view( + device_ids.data(), static_cast(size)); + auto built = build_from_device_ids(res, device_ids_view, pre_sorted); + return roaring_allowlist{std::make_unique(res, dataset_rows, std::move(built))}; +} + +roaring_allowlist roaring_allowlist::from_ids( + raft::resources const& res, + std::size_t dataset_rows, + raft::device_vector_view ids, + bool pre_sorted) +{ + validate_dataset_rows(dataset_rows); + auto built = build_from_device_ids(res, ids, pre_sorted); + return roaring_allowlist{std::make_unique(res, dataset_rows, std::move(built))}; +} + +roaring_allowlist::~roaring_allowlist() = default; +roaring_allowlist::roaring_allowlist(roaring_allowlist&&) noexcept = default; +roaring_allowlist& roaring_allowlist::operator=(roaring_allowlist&&) noexcept = default; + +std::size_t roaring_allowlist::dataset_rows() const noexcept { return impl_->dataset_rows_; } + +std::size_t roaring_allowlist::cardinality() const noexcept { return impl_->cardinality_; } + +bool roaring_allowlist::empty() const noexcept { return cardinality() == 0; } + +std::size_t roaring_allowlist::size_bytes() const noexcept +{ + auto bytes = impl_->device_reference.size() * sizeof(ref_type); + if (impl_->owner) { bytes += static_cast(impl_->owner->size_bytes()); } + return bytes; +} + +roaring_allowlist_view roaring_allowlist::view() const noexcept +{ + return roaring_allowlist_view{impl_->reference(), dataset_rows(), cardinality()}; +} + +} // namespace cuvs::core diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index b23b264dcb..15cd4d3418 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -441,6 +442,25 @@ void search(raft::resources const& res, } catch (const std::bad_cast&) { } + try { + auto& sample_filter = + dynamic_cast(sample_filter_ref); + RAFT_EXPECTS(sample_filter.valid(), "roaring_filter must be initialized before search."); + RAFT_EXPECTS(sample_filter.num_queries() == static_cast(queries.extent(0)), + "Roaring filter query rows must equal the number of search queries."); + RAFT_EXPECTS(sample_filter.dataset_rows() == static_cast(idx.dataset().n_rows()), + "Roaring filter dataset_rows must equal the number of rows in the index."); + + search_params params_copy = params; + if (params.filtering_rate < 0.0f) { + params_copy.filtering_rate = sample_filter.filtering_rate(); + } + auto sample_filter_copy = sample_filter; + return search_with_filtering( + res, params_copy, idx, queries, neighbors, distances, sample_filter_copy); + } catch (const std::bad_cast&) { + } + try { auto& sample_filter = dynamic_cast(sample_filter_ref); diff --git a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp index e0406f5d1a..f674bf3f3c 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp @@ -146,6 +146,12 @@ struct is_bloom_filter : std::false_type {}; template <> struct is_bloom_filter<::cuvs::neighbors::filtering::bloom_filter> : std::true_type {}; +template +struct is_roaring_filter : std::false_type {}; + +template <> +struct is_roaring_filter<::cuvs::neighbors::filtering::roaring_filter> : std::true_type {}; + template struct is_udf_filter : std::false_type {}; @@ -198,6 +204,8 @@ void fill_cagra_sample_filter(cagra_sample_filter& out, out.filter_data = get_cagra_device_payload(make_cagra_bloom_filter_storage(filter), stream); } else if constexpr (is_udf_filter::value) { out.filter_data = filter.filter_data; + } else if constexpr (is_roaring_filter::value) { + out.filter_data = filter.device_payload(); } } @@ -211,6 +219,8 @@ std::uint64_t cagra_filter_payload_hash(const FilterT& filter) return cagra_payload_hash(make_cagra_bloom_filter_storage(filter)); } else if constexpr (requires { filter.filter; }) { return cagra_filter_payload_hash(filter.filter); + } else if constexpr (is_roaring_filter::value) { + return cagra_payload_hash(filter.device_payload()); } else { return 0; } @@ -222,6 +232,8 @@ void* cagra_filter_data_ptr(const FilterT& filter) using DecayedFilter = std::decay_t; if constexpr (is_bloom_filter::value || is_udf_filter::value) { return filter.filter_data; + } else if constexpr (is_roaring_filter::value) { + return filter.device_payload(); } else if constexpr (requires { filter.filter; }) { return cagra_filter_data_ptr(filter.filter); } else { diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index a935df9b6f..cf15999842 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -77,6 +77,8 @@ cuvs::neighbors::cagra::index merge_rebuild( "Bitmap filter isn't supported inside cagra::merge"); RAFT_EXPECTS(row_filter.get_filter_type() != cuvs::neighbors::filtering::FilterType::Bloom, "Bloom filter isn't supported inside cagra::merge"); + RAFT_EXPECTS(row_filter.get_filter_type() != cuvs::neighbors::filtering::FilterType::Roaring, + "Roaring filter isn't supported inside cagra::merge"); for (cagra_index_t* index : indices) { RAFT_EXPECTS(index != nullptr, diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_impl.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_impl.cuh index 1b3b3825f2..ac94b119ab 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_impl.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_impl.cuh @@ -7,6 +7,7 @@ #include "extern_device_functions.cuh" +#include "../../roaring_filter_data.cuh" #include "../../sample_filter_data.cuh" #include @@ -51,4 +52,19 @@ __device__ bool sample_filter_bloom_filter_impl(uint32_t /*query_id*/, return data->filter.contains(static_cast(node_id)); } +template +__device__ bool sample_filter_roaring_impl(uint32_t query_id, + SourceIndexT node_id, + void* filter_data) +{ + if (filter_data == nullptr) { return false; } + + auto const* data = static_cast const*>(filter_data); + if (query_id >= data->num_queries || static_cast(node_id) >= data->dataset_rows || + data->empty_rows[query_id] != 0) { + return false; + } + return data->refs[query_id]->contains(static_cast(node_id)); +} + } // namespace cuvs::neighbors::detail diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_matrix.json b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_matrix.json index b58f56ceb6..7ce27397ff 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_matrix.json +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_matrix.json @@ -1,5 +1,5 @@ { - "filter_name": ["none", "bitset", "bloom_filter"], + "filter_name": ["none", "bitset", "bloom_filter", "roaring"], "_bitset": [ { "bitset_type": "uint32_t", diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in b/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in index d53dabf248..759d66a1a9 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in @@ -13,6 +13,8 @@ using bitset_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQue cuvs::neighbors::filtering::bitset_filter>; using bloom_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< cuvs::neighbors::filtering::bloom_filter>; +using roaring_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< + cuvs::neighbors::filtering::roaring_filter>; using udf_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< cuvs::neighbors::filtering::udf_filter>; @@ -25,6 +27,7 @@ instantiate_kernel_selection(data_t, cuvs::neighbors::filtering::none_sample_filter); instantiate_kernel_selection(data_t, uint32_t, float, bitset_filter_t); instantiate_kernel_selection(data_t, uint32_t, float, bloom_filter_t); +instantiate_kernel_selection(data_t, uint32_t, float, roaring_filter_t); instantiate_kernel_selection(data_t, uint32_t, float, udf_filter_t); instantiate_kernel_selection_mp(data_t, uint32_t, diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in b/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in index bf939b9e25..530f61e9e6 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in @@ -13,6 +13,8 @@ using bitset_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQue cuvs::neighbors::filtering::bitset_filter>; using bloom_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< cuvs::neighbors::filtering::bloom_filter>; +using roaring_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< + cuvs::neighbors::filtering::roaring_filter>; using udf_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< cuvs::neighbors::filtering::udf_filter>; @@ -25,6 +27,7 @@ instantiate_kernel_selection(data_t, cuvs::neighbors::filtering::none_sample_filter); instantiate_kernel_selection(data_t, uint32_t, float, bitset_filter_t); instantiate_kernel_selection(data_t, uint32_t, float, bloom_filter_t); +instantiate_kernel_selection(data_t, uint32_t, float, roaring_filter_t); instantiate_kernel_selection(data_t, uint32_t, float, udf_filter_t); instantiate_kernel_selection_mp(data_t, uint32_t, diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh index 18dd5345e4..cd5dee24ef 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh @@ -84,6 +84,8 @@ std::uint64_t cagra_sample_filter_type_id(const SampleFilterT& sample_filter) return 3; } else if constexpr (is_bloom_filter::value) { return 2; + } else if constexpr (is_roaring_filter::value) { + return 4; } else if constexpr (is_bitset_filter::value) { return 1; } else if constexpr (requires { sample_filter.filter; }) { diff --git a/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp b/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp index 72797ec0be..c563759db4 100644 --- a/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp +++ b/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp @@ -103,6 +103,8 @@ struct sample_filter_jit_tag { return cuvs::neighbors::detail::tag_filter_none{}; } else if constexpr (is_bloom_filter::value) { return cuvs::neighbors::detail::tag_filter_bloom_filter{}; + } else if constexpr (is_roaring_filter::value) { + return cuvs::neighbors::detail::tag_filter_roaring{}; } else if constexpr (is_udf_filter::value) { return cuvs::neighbors::detail::tag_filter_udf{}; } else if constexpr (requires { std::declval().filter; }) { @@ -114,6 +116,8 @@ struct sample_filter_jit_tag { return cuvs::neighbors::detail::tag_filter_bitset{}; } else if constexpr (is_bloom_filter>::value) { return cuvs::neighbors::detail::tag_filter_bloom_filter{}; + } else if constexpr (is_roaring_filter>::value) { + return cuvs::neighbors::detail::tag_filter_roaring{}; } else if constexpr (is_udf_filter>::value) { return cuvs::neighbors::detail::tag_filter_udf{}; } else { diff --git a/cpp/src/neighbors/detail/dynamic_batching.cuh b/cpp/src/neighbors/detail/dynamic_batching.cuh index 983b021f55..2b8ee12206 100644 --- a/cpp/src/neighbors/detail/dynamic_batching.cuh +++ b/cpp/src/neighbors/detail/dynamic_batching.cuh @@ -843,6 +843,16 @@ RAFT_KERNEL scatter_outputs( * The search function must be thread-safe. We only have to pay attention to the `mutable` members * though, because the function is marked const. */ +inline auto validate_sample_filter(cuvs::neighbors::filtering::base_filter const* sample_filter) + -> cuvs::neighbors::filtering::base_filter const* +{ + RAFT_EXPECTS( + sample_filter == nullptr || + sample_filter->get_filter_type() != cuvs::neighbors::filtering::FilterType::Roaring, + "dynamic_batching does not support roaring_filter; use direct cagra::search instead."); + return sample_filter; +} + template class batch_runner { public: @@ -860,7 +870,10 @@ class batch_runner { upstream_search_type_const* upstream_search, const cuvs::neighbors::filtering::base_filter* sample_filter) : res_{res}, - upstream_search_{[&upstream_index, upstream_search, upstream_params, sample_filter]( + upstream_search_{[&upstream_index, + upstream_search, + upstream_params, + sample_filter = validate_sample_filter(sample_filter)]( raft::resources const& res, raft::device_matrix_view queries, raft::device_matrix_view neighbors, diff --git a/cpp/src/neighbors/detail/roaring_filter_data.cuh b/cpp/src/neighbors/detail/roaring_filter_data.cuh new file mode 100644 index 0000000000..8b908ee2dc --- /dev/null +++ b/cpp/src/neighbors/detail/roaring_filter_data.cuh @@ -0,0 +1,27 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include + +namespace cuvs::neighbors::detail { + +/// Per-query cuco Roaring bitmap references for linked @c sample_filter in CAGRA JIT LTO. +template +struct roaring_filter_data_t { + using ref_type = cuco::experimental::roaring_bitmap_ref; + + // Each entry points at a reference owned and initialized by one roaring_allowlist. The pointer + // table is built by roaring_filter; CAGRA only follows it. + ref_type const* const* refs{nullptr}; + std::uint8_t const* empty_rows{nullptr}; + std::uint32_t num_queries{}; + std::uint64_t dataset_rows{}; +}; + +} // namespace cuvs::neighbors::detail diff --git a/cpp/src/neighbors/detail/tiered_index.cuh b/cpp/src/neighbors/detail/tiered_index.cuh index 8db27d7d6c..3ee63b2b04 100644 --- a/cpp/src/neighbors/detail/tiered_index.cuh +++ b/cpp/src/neighbors/detail/tiered_index.cuh @@ -183,6 +183,10 @@ struct index_state { raft::device_matrix_view distances, const cuvs::neighbors::filtering::base_filter& sample_filter) { + RAFT_EXPECTS( + sample_filter.get_filter_type() != cuvs::neighbors::filtering::FilterType::Roaring, + "tiered_index::search does not support roaring_filter; use direct cagra::search instead."); + // if we only have ANN vectors, search those and return immendiately if (bfknn_rows() == 0) { search_fn(res, search_params, *ann_index, queries, neighbors, distances, sample_filter); diff --git a/cpp/src/neighbors/roaring_filter.cu b/cpp/src/neighbors/roaring_filter.cu new file mode 100644 index 0000000000..bfd56a6948 --- /dev/null +++ b/cpp/src/neighbors/roaring_filter.cu @@ -0,0 +1,179 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "detail/roaring_filter_data.cuh" + +#include +#include + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::filtering { +namespace { + +using data_type = cuvs::neighbors::detail::roaring_filter_data_t; +using ref_type = data_type::ref_type; + +std::size_t validate_views(std::span allowlists) +{ + RAFT_EXPECTS(!allowlists.empty(), "roaring_filter requires at least one query allowlist."); + RAFT_EXPECTS(allowlists.front().valid(), "roaring_filter requires valid allowlist views."); + auto const dataset_rows = allowlists.front().dataset_rows(); + for (auto const& allowlist : allowlists) { + RAFT_EXPECTS(allowlist.valid(), "roaring_filter requires valid allowlist views."); + RAFT_EXPECTS(allowlist.dataset_rows() == dataset_rows, + "Every roaring_filter allowlist must have the same dataset_rows."); + RAFT_EXPECTS(allowlist.empty() || allowlist.device_reference() != nullptr, + "A nonempty Roaring allowlist must carry a device reference."); + } + RAFT_EXPECTS(allowlists.size() <= std::numeric_limits::max(), + "roaring_filter has too many query allowlists."); + return dataset_rows; +} + +float estimate_filtering_rate(std::span allowlists, + std::size_t dataset_rows) +{ + // CAGRA accepts one filtering-rate hint for the entire batch. Use the sparsest allowlist so no + // query is under-provisioned; callers can override the hint when throughput is more important. + auto minimum_cardinality = dataset_rows; + for (auto const& allowlist : allowlists) { + minimum_cardinality = std::min(minimum_cardinality, allowlist.cardinality()); + } + auto const rejected = + 1.0 - static_cast(minimum_cardinality) / static_cast(dataset_rows); + return std::clamp(static_cast(rejected), 0.0f, 0.999f); +} + +} // namespace + +struct roaring_filter::impl { + std::vector allowlists; + rmm::device_uvector refs; + rmm::device_uvector empty_rows; + rmm::device_uvector payload; + std::size_t dataset_rows_{}; + float filtering_rate_{}; + + impl(raft::resources const& res, + std::span input_allowlists) + : allowlists(input_allowlists.begin(), input_allowlists.end()), + refs(input_allowlists.size(), raft::resource::get_cuda_stream(res)), + empty_rows(input_allowlists.size(), raft::resource::get_cuda_stream(res)), + payload(1, raft::resource::get_cuda_stream(res)), + dataset_rows_(validate_views(input_allowlists)), + filtering_rate_(estimate_filtering_rate(input_allowlists, dataset_rows_)) + { + auto stream = raft::resource::get_cuda_stream(res); + std::vector host_refs; + std::vector host_empty; + host_refs.reserve(allowlists.size()); + host_empty.reserve(allowlists.size()); + for (auto const& allowlist : allowlists) { + host_refs.push_back(static_cast(allowlist.device_reference())); + host_empty.push_back(allowlist.empty() ? 1 : 0); + } + + raft::copy(refs.data(), host_refs.data(), host_refs.size(), stream); + raft::copy(empty_rows.data(), host_empty.data(), host_empty.size(), stream); + auto const host_payload = data_type{refs.data(), + empty_rows.data(), + static_cast(allowlists.size()), + static_cast(dataset_rows_)}; + raft::copy(payload.data(), &host_payload, 1, stream); + + // Construction establishes a stream-independent ready object. Search only reads these stable + // allocations and therefore needs no event, copy, initialization kernel, or synchronization. + raft::resource::sync_stream(res); + } + + void recompute_filtering_rate() + { + filtering_rate_ = estimate_filtering_rate(allowlists, dataset_rows_); + } +}; + +roaring_filter::roaring_filter(raft::resources const& res, + std::span allowlists) + : impl_(std::make_shared(res, allowlists)) +{ +} + +bool roaring_filter::valid() const noexcept { return impl_ != nullptr; } + +std::size_t roaring_filter::num_queries() const noexcept +{ + return valid() ? impl_->allowlists.size() : 0; +} + +std::size_t roaring_filter::dataset_rows() const noexcept +{ + return valid() ? impl_->dataset_rows_ : 0; +} + +std::size_t roaring_filter::cardinality(std::size_t query_id) const +{ + RAFT_EXPECTS(valid(), "roaring_filter is not initialized."); + RAFT_EXPECTS(query_id < num_queries(), "roaring_filter query_id is out of range."); + return impl_->allowlists[query_id].cardinality(); +} + +bool roaring_filter::empty(std::size_t query_id) const { return cardinality(query_id) == 0; } + +float roaring_filter::filtering_rate() const noexcept +{ + return valid() ? impl_->filtering_rate_ : 0.0f; +} + +std::size_t roaring_filter::size_bytes() const noexcept +{ + if (!valid()) { return 0; } + return impl_->refs.size() * sizeof(ref_type const*) + + impl_->empty_rows.size() * sizeof(std::uint8_t) + + impl_->payload.size() * sizeof(data_type); +} + +void roaring_filter::set_allowlist(raft::resources const& res, + std::size_t query_id, + cuvs::core::roaring_allowlist_view replacement) +{ + RAFT_EXPECTS(valid(), "roaring_filter is not initialized."); + RAFT_EXPECTS(query_id < num_queries(), "roaring_filter query_id is out of range."); + RAFT_EXPECTS(replacement.valid(), "roaring_filter requires a valid replacement allowlist view."); + RAFT_EXPECTS(replacement.dataset_rows() == dataset_rows(), + "Replacement Roaring allowlist must have the filter's dataset_rows."); + RAFT_EXPECTS(replacement.empty() || replacement.device_reference() != nullptr, + "A nonempty replacement Roaring allowlist must carry a device reference."); + + auto stream = raft::resource::get_cuda_stream(res); + auto const ref = static_cast(replacement.device_reference()); + auto const empty_ = static_cast(replacement.empty() ? 1 : 0); + raft::copy(impl_->refs.data() + query_id, &ref, 1, stream); + raft::copy(impl_->empty_rows.data() + query_id, &empty_, 1, stream); + raft::resource::sync_stream(res); + + impl_->allowlists[query_id] = replacement; + impl_->recompute_filtering_rate(); +} + +void* roaring_filter::device_payload() const noexcept +{ + return valid() ? const_cast(impl_->payload.data()) : nullptr; +} + +} // namespace cuvs::neighbors::filtering diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 7d3720be08..1a41bbc570 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -204,6 +204,14 @@ ConfigureTest( PERCENT 100 ) +ConfigureTest( + NAME CORE_ROARING_ALLOWLIST_TEST + PATH neighbors/roaring_allowlist.cu + ADDITIONAL_DEP cuco::cuco + GPUS 1 + PERCENT 100 +) + ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_TEST_BUGS PATH neighbors/ann_cagra/bug_extreme_inputs_oob.cu diff --git a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu index 2cc92b2aca..e510b11709 100644 --- a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu +++ b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu @@ -8,6 +8,7 @@ #include "../ann_cagra.cuh" #include +#include #include #include @@ -19,6 +20,8 @@ #include #include +#include +#include #include #include #include @@ -348,7 +351,8 @@ TEST_P(CagraUdfFilterTest, TenantContextHonorsQuerySpecificMetadata) std::vector host_row_tenants(n_rows); std::vector host_query_tenants(n_queries); for (int64_t i = 0; i < n_rows; ++i) { - host_row_tenants[static_cast(i)] = static_cast((i / 5) % 3); + // Equal tenant cardinalities make the query-specific reference comparison deterministic. + host_row_tenants[static_cast(i)] = static_cast(i % 3); } for (int64_t q = 0; q < n_queries; ++q) { host_query_tenants[static_cast(q)] = static_cast(q % 3); @@ -379,6 +383,74 @@ TEST_P(CagraUdfFilterTest, TenantContextHonorsQuerySpecificMetadata) EXPECT_EQ(host_row_tenants[source_id], query_tenant); } } + + // Build independent reusable owners and map one view to each query. Compare Roaring against this + // existing query-specific UDF result for every single-partition CAGRA algorithm. max_queries=2 + // above also verifies query-offset propagation through internal chunking. + std::vector tenant_allowlists; + tenant_allowlists.reserve(n_queries); + for (std::int64_t query = 0; query < n_queries; ++query) { + std::vector allowed_ids; + auto const query_tenant = host_query_tenants[static_cast(query)]; + for (std::int64_t row = 0; row < n_rows; ++row) { + if (host_row_tenants[static_cast(row)] == query_tenant) { + allowed_ids.push_back(static_cast(row)); + } + } + tenant_allowlists.push_back(cuvs::core::roaring_allowlist::from_ids( + res, + n_rows, + raft::make_host_vector_view(allowed_ids.data(), + allowed_ids.size()), + true)); + } + + std::vector tenant_views; + tenant_views.reserve(n_queries); + for (auto const& allowlist : tenant_allowlists) { + tenant_views.push_back(allowlist.view()); + } + cuvs::neighbors::filtering::roaring_filter roaring_filter(res, tenant_views); + auto roaring_result = search(roaring_filter, 2.0f / 3.0f); + expect_same_results(result, roaring_result); + + // Exercise both device tables: first mark one query empty, then repoint that same slot to a + // different reusable owner and verify CAGRA consumes the updated mapping. + auto empty_allowlist = cuvs::core::roaring_allowlist::from_ids( + res, n_rows, raft::make_host_vector_view(nullptr, 0)); + roaring_filter.set_allowlist(res, 1, empty_allowlist.view()); + auto empty_result = search(roaring_filter, 0.999f); + for (std::int64_t i = 0; i < k; ++i) { + auto const source_id = empty_result.neighbors[static_cast(k + i)]; + EXPECT_GE(source_id, static_cast(n_rows)); + } + + roaring_filter.set_allowlist(res, 1, tenant_allowlists.front().view()); + auto updated_result = search(roaring_filter, 2.0f / 3.0f); + for (std::int64_t query = 0; query < n_queries; ++query) { + auto const expected_tenant = + query == 1 ? std::uint32_t{0} : host_query_tenants[static_cast(query)]; + for (std::int64_t i = 0; i < k; ++i) { + auto const source_id = updated_result.neighbors[static_cast(query * k + i)]; + ASSERT_LT(source_id, static_cast(n_rows)); + EXPECT_EQ(host_row_tenants[source_id], expected_tenant); + } + } + + if (GetParam() == cagra::search_algo::SINGLE_CTA) { + auto wrong_queries = cuvs::core::roaring_allowlist::from_ids( + res, n_rows, raft::make_host_vector_view(nullptr, 0)); + std::array wrong_query_views{wrong_queries.view()}; + cuvs::neighbors::filtering::roaring_filter wrong_query_filter(res, wrong_query_views); + EXPECT_THROW(search(wrong_query_filter), raft::logic_error); + + auto wrong_columns = cuvs::core::roaring_allowlist::from_ids( + res, n_rows + 1, raft::make_host_vector_view(nullptr, 0)); + std::vector wrong_column_views( + static_cast(n_queries), wrong_columns.view()); + cuvs::neighbors::filtering::roaring_filter wrong_column_filter(res, wrong_column_views); + EXPECT_THROW(search(wrong_column_filter), raft::logic_error); + } } INSTANTIATE_TEST_CASE_P(CagraUdfFilters, diff --git a/cpp/tests/neighbors/roaring_allowlist.cu b/cpp/tests/neighbors/roaring_allowlist.cu new file mode 100644 index 0000000000..5e07416926 --- /dev/null +++ b/cpp/tests/neighbors/roaring_allowlist.cu @@ -0,0 +1,226 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::core { +namespace { + +roaring_allowlist from_ids(raft::resources const& res, + std::size_t dataset_rows, + std::vector const& ids, + bool pre_sorted = false) +{ + return roaring_allowlist::from_ids( + res, + dataset_rows, + raft::make_host_vector_view(ids.data(), ids.size()), + pre_sorted); +} + +roaring_allowlist from_device_ids(raft::resources const& res, + std::size_t dataset_rows, + std::vector const& ids, + bool pre_sorted = false) +{ + auto device_ids = raft::make_device_vector(res, ids.size()); + auto const stream = raft::resource::get_cuda_stream(res); + raft::update_device(device_ids.data_handle(), ids.data(), ids.size(), stream); + return roaring_allowlist::from_ids( + res, + dataset_rows, + raft::make_device_vector_view(device_ids.data_handle(), + ids.size()), + pre_sorted); +} + +using ref_type = cuco::experimental::roaring_bitmap_ref; + +__global__ void membership_probe_kernel(ref_type const* reference, + bool empty, + std::uint32_t const* row_ids, + std::uint8_t* output, + std::size_t size) +{ + auto const i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < size) { output[i] = !empty && reference->contains(row_ids[i]); } +} + +void expect_membership(raft::resources const& res, + roaring_allowlist const& allowlist, + std::vector const& row_ids, + std::vector const& expected) +{ + ASSERT_EQ(row_ids.size(), expected.size()); + auto rows = raft::make_device_vector(res, row_ids.size()); + auto output = raft::make_device_vector(res, expected.size()); + auto stream = raft::resource::get_cuda_stream(res); + raft::update_device(rows.data_handle(), row_ids.data(), row_ids.size(), stream); + constexpr std::size_t block_size = 256; + raft::launch_kernel(stream, + dim3((row_ids.size() + block_size - 1) / block_size), + dim3(block_size), + membership_probe_kernel, + static_cast(allowlist.view().device_reference()), + allowlist.empty(), + rows.data_handle(), + output.data_handle(), + row_ids.size()); + + std::vector actual(expected.size()); + raft::update_host(actual.data(), output.data_handle(), actual.size(), stream); + raft::resource::sync_stream(res); + EXPECT_EQ(actual, expected); +} + +TEST(RoaringAllowlist, BuildsArrayBitmapAndMultipleContainers) +{ + raft::device_resources res; + + std::vector ids{65537, 7, 3, 131074, 5}; + auto const original = ids; + auto array = from_ids(res, 200000, ids); + EXPECT_EQ(ids, original); + EXPECT_EQ(array.dataset_rows(), 200000); + EXPECT_EQ(array.cardinality(), 5); + EXPECT_FALSE(array.empty()); + EXPECT_GT(array.size_bytes(), 0); + EXPECT_TRUE(array.view().valid()); + EXPECT_EQ(array.view().cardinality(), 5); + expect_membership(res, array, {3, 4, 7, 65537, 131073, 131074, 200000}, {1, 0, 1, 1, 0, 1, 0}); + + std::vector dense; + for (std::uint32_t id = 0; id < 10000; id += 2) { + dense.push_back(id); + } + auto bitmap = from_ids(res, 10000, dense, true); + EXPECT_EQ(bitmap.cardinality(), 5000); + expect_membership(res, bitmap, {0, 1, 4096, 9998, 9999}, {1, 0, 1, 1, 0}); +} + +TEST(RoaringAllowlist, SupportsPresortedHostAndDeviceInputs) +{ + raft::device_resources res; + std::vector const ids{1, 4, 7, 65536, 65539}; + auto host = from_ids(res, 131072, ids, true); + auto device = from_device_ids(res, 131072, ids, true); + + EXPECT_EQ(host.cardinality(), ids.size()); + EXPECT_EQ(device.cardinality(), ids.size()); + expect_membership(res, host, {0, 1, 7, 8, 65539}, {0, 1, 1, 0, 1}); + expect_membership(res, device, {0, 1, 7, 8, 65539}, {0, 1, 1, 0, 1}); +} + +TEST(RoaringAllowlist, HandlesEmptyAndFullUint32Domain) +{ + raft::device_resources res; + + auto empty = from_ids(res, 32, {}); + EXPECT_TRUE(empty.empty()); + EXPECT_EQ(empty.cardinality(), 0); + EXPECT_TRUE(empty.view().valid()); + EXPECT_EQ(empty.view().device_reference(), nullptr); + expect_membership(res, empty, {0, 31, 32}, {0, 0, 0}); + + EXPECT_THROW(from_ids(res, 0, {}), raft::logic_error); + + auto maximum = + from_ids(res, std::uint64_t{1} << 32, {0, std::numeric_limits::max()}, true); + expect_membership(res, maximum, {0, 1, std::numeric_limits::max()}, {1, 0, 1}); +} + +TEST(RoaringAllowlist, ViewIsZeroCopyAndSurvivesOwnerMove) +{ + raft::device_resources res; + auto allowlist = from_ids(res, 32, {1, 4, 7}); + auto first_view = allowlist.view(); + auto next_view = allowlist.view(); + EXPECT_EQ(first_view.device_reference(), next_view.device_reference()); + + auto moved = std::move(allowlist); + auto moved_view = moved.view(); + EXPECT_EQ(first_view.device_reference(), moved_view.device_reference()); + EXPECT_EQ(moved_view.cardinality(), 3); + expect_membership(res, moved, {1, 2, 7}, {1, 0, 1}); +} + +TEST(RoaringAllowlist, StreamOrderedConstructionSupportsEventHandoff) +{ + rmm::cuda_stream build_stream; + rmm::cuda_stream consume_stream; + raft::device_resources build_res; + raft::device_resources consume_res; + raft::resource::set_cuda_stream(build_res, build_stream.view()); + raft::resource::set_cuda_stream(consume_res, consume_stream.view()); + + auto allowlist = from_ids(build_res, 1000, {900, 100, 300, 200}); + + cudaEvent_t ready{}; + RAFT_CUDA_TRY(cudaEventCreateWithFlags(&ready, cudaEventDisableTiming)); + RAFT_CUDA_TRY(cudaEventRecord(ready, build_stream.value())); + RAFT_CUDA_TRY(cudaStreamWaitEvent(consume_stream.value(), ready)); + expect_membership(consume_res, allowlist, {99, 100, 200, 300, 900}, {0, 1, 1, 1, 1}); + RAFT_CUDA_TRY(cudaEventDestroy(ready)); +} + +TEST(RoaringFilter, ReusesViewsUpdatesMappingsAndRejectsInvalidInputs) +{ + raft::device_resources res; + auto first = from_ids(res, 16, {1, 3}); + auto second = from_ids(res, 16, {2, 4, 6}); + auto empty = from_ids(res, 16, {}); + + std::array views{first.view(), second.view(), first.view()}; + cuvs::neighbors::filtering::roaring_filter filter(res, views); + EXPECT_TRUE(filter.valid()); + EXPECT_EQ(filter.num_queries(), 3); + EXPECT_EQ(filter.dataset_rows(), 16); + EXPECT_EQ(filter.cardinality(0), 2); + EXPECT_EQ(filter.cardinality(1), 3); + // The automatic batch hint follows the sparsest row: 1 - 2 / 16. + EXPECT_FLOAT_EQ(filter.filtering_rate(), 0.875f); + + auto const* payload = filter.device_payload(); + auto shared_copy = filter; + shared_copy.set_allowlist(res, 1, empty.view()); + EXPECT_EQ(filter.device_payload(), payload); + EXPECT_TRUE(filter.empty(1)); + EXPECT_FLOAT_EQ(filter.filtering_rate(), 0.999f); + + EXPECT_THROW(cuvs::neighbors::filtering::roaring_filter( + res, std::span{}), + raft::logic_error); + auto different_shape = from_ids(res, 17, {1}); + std::array mismatched{first.view(), different_shape.view()}; + EXPECT_THROW(cuvs::neighbors::filtering::roaring_filter(res, mismatched), raft::logic_error); + EXPECT_THROW(filter.set_allowlist(res, 3, first.view()), raft::logic_error); + EXPECT_THROW(filter.set_allowlist(res, 0, different_shape.view()), raft::logic_error); +} + +} // namespace +} // namespace cuvs::core diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index f4dffad3a9..d4edf4d127 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -41,6 +41,7 @@ add_executable(BALANCED_KMEANS_EXAMPLE src/balanced_kmeans_example.cu) add_executable(CAGRA_EXAMPLE src/cagra_example.cu) add_executable(CAGRA_FILTER_UDF_EXAMPLE src/cagra_filter_udf_example.cu) add_executable(CAGRA_BLOOM_FILTER_EXAMPLE src/cagra_bloom_filter_example.cu) +add_executable(CAGRA_ROARING_FILTER_EXAMPLE src/cagra_roaring_filter_example.cu) add_executable(CAGRA_HNSW_ACE_BUILD_EXAMPLE src/cagra_hnsw_ace_build.cu) add_executable(CAGRA_HNSW_ACE_EXAMPLE src/cagra_hnsw_ace_example.cu) add_executable(CAGRA_PERSISTENT_EXAMPLE src/cagra_persistent_example.cu) @@ -63,6 +64,9 @@ target_link_libraries( target_link_libraries( CAGRA_BLOOM_FILTER_EXAMPLE PRIVATE cuvs::cuvs cuco::cuco $ ) +target_link_libraries( + CAGRA_ROARING_FILTER_EXAMPLE PRIVATE cuvs::cuvs $ +) target_link_libraries( CAGRA_HNSW_ACE_BUILD_EXAMPLE PRIVATE cuvs::cuvs $ ) diff --git a/examples/cpp/src/cagra_roaring_filter_example.cu b/examples/cpp/src/cagra_roaring_filter_example.cu new file mode 100644 index 0000000000..4835af001f --- /dev/null +++ b/examples/cpp/src/cagra_roaring_filter_example.cu @@ -0,0 +1,107 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace { + +constexpr std::int64_t n_rows = 1024; +constexpr std::int64_t n_dim = 16; +constexpr std::int64_t n_queries = 4; +constexpr std::int64_t k = 8; + +} // namespace + +int main() +{ + raft::device_resources res; + rmm::mr::pool_memory_resource pool_mr(rmm::mr::get_current_device_resource_ref(), + 1024 * 1024 * 1024ull); + rmm::mr::set_current_device_resource(pool_mr); + + auto dataset = raft::make_device_matrix(res, n_rows, n_dim); + auto queries = raft::make_device_matrix(res, n_queries, n_dim); + raft::random::RngState rng(1234ULL); + raft::random::uniform(res, rng, dataset.data_handle(), dataset.size(), -1.0f, 1.0f); + raft::random::uniform(res, rng, queries.data_handle(), queries.size(), -1.0f, 1.0f); + + cuvs::neighbors::cagra::index_params index_params; + index_params.graph_degree = 32; + index_params.intermediate_graph_degree = 64; + auto padded = + cuvs::neighbors::make_device_padded_dataset_view(res, raft::make_const_mdspan(dataset.view())); + auto index = cuvs::neighbors::cagra::build(res, index_params, padded); + + // Each owner is independently reusable. The filter supplies the query-to-allowlist mapping by + // retaining only zero-copy views of their already initialized device references. + std::vector owners; + owners.reserve(n_queries); + for (std::uint32_t query = 0; query < n_queries; ++query) { + std::vector ids; + for (std::uint32_t row = query; row < n_rows; row += n_queries) { + ids.push_back(row); + } + owners.push_back(cuvs::core::roaring_allowlist::from_ids( + res, + n_rows, + raft::make_host_vector_view(ids.data(), ids.size()), + true)); + } + + std::vector views; + views.reserve(owners.size()); + for (auto const& owner : owners) { + views.push_back(owner.view()); + } + cuvs::neighbors::filtering::roaring_filter filter(res, views); + auto const* prepared_payload = filter.device_payload(); + + auto neighbors = raft::make_device_matrix(res, n_queries, k); + auto distances = raft::make_device_matrix(res, n_queries, k); + cuvs::neighbors::cagra::search_params search_params; + search_params.algo = cuvs::neighbors::cagra::search_algo::MULTI_CTA; + search_params.itopk_size = 64; + search_params.max_queries = 2; // also demonstrates internal query chunking + + cuvs::neighbors::cagra::search(res, + search_params, + index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view(), + filter); + if (filter.device_payload() != prepared_payload) { return 1; } + + std::vector host_neighbors(neighbors.size()); + raft::copy(host_neighbors.data(), + neighbors.data_handle(), + host_neighbors.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (std::int64_t query = 0; query < n_queries; ++query) { + for (std::int64_t rank = 0; rank < k; ++rank) { + auto row = host_neighbors[static_cast(query * k + rank)]; + if (row >= n_rows || row % n_queries != static_cast(query)) { return 1; } + } + } + + std::cout << "CAGRA reused " << owners.size() + << " independently owned Roaring allowlists through one prepared filter payload.\\n"; + return 0; +}