From f6406773f834d8c34ac0478a85c7424abdf25ecf Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Sat, 15 Aug 2026 00:10:43 +0000 Subject: [PATCH 1/5] Initial impl --- cpp/CMakeLists.txt | 2 + cpp/include/cuvs/core/roaring_allowlist.hpp | 174 ++ .../cuvs/detail/jit_lto/common_fragments.hpp | 1 + cpp/include/cuvs/neighbors/common.hpp | 92 +- cpp/src/core/roaring_allowlist.cu | 2249 +++++++++++++++++ cpp/src/neighbors/cagra.cuh | 20 + .../detail/cagra/cagra_filter_payload.hpp | 12 + .../neighbors/detail/cagra/cagra_merge.cuh | 2 + .../jit_lto_kernels/sample_filter_impl.cuh | 16 + .../jit_lto_kernels/sample_filter_matrix.json | 2 +- .../detail/cagra/search_multi_cta_inst.cu.in | 3 + .../detail/cagra/search_single_cta_inst.cu.in | 3 + .../search_single_cta_kernel_launcher_jit.cuh | 2 + .../detail/cagra/shared_launcher_jit.hpp | 4 + .../neighbors/detail/roaring_filter_data.cuh | 27 + cpp/src/neighbors/roaring_filter.cu | 190 ++ cpp/tests/CMakeLists.txt | 8 + .../neighbors/ann_cagra/test_filter_udf.cu | 115 +- cpp/tests/neighbors/roaring_allowlist.cu | 944 +++++++ examples/cpp/CMakeLists.txt | 4 + .../cpp/src/cagra_roaring_filter_example.cu | 174 ++ 21 files changed, 4040 insertions(+), 4 deletions(-) create mode 100644 cpp/include/cuvs/core/roaring_allowlist.hpp create mode 100644 cpp/src/core/roaring_allowlist.cu create mode 100644 cpp/src/neighbors/detail/roaring_filter_data.cuh create mode 100644 cpp/src/neighbors/roaring_filter.cu create mode 100644 cpp/tests/neighbors/roaring_allowlist.cu create mode 100644 examples/cpp/src/cagra_roaring_filter_example.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 088c5f689c..933607c283 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1358,6 +1358,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/include/cuvs/core/roaring_allowlist.hpp b/cpp/include/cuvs/core/roaring_allowlist.hpp new file mode 100644 index 0000000000..af190454f9 --- /dev/null +++ b/cpp/include/cuvs/core/roaring_allowlist.hpp @@ -0,0 +1,174 @@ +/* + * 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 row in a batched 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 pre-parsed 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 batch of exact per-query Roaring allowlists. + * + * Logically, the owner is a sparse matrix with one allowlist row per query and one possible column + * per dataset row. @ref from_ids accepts one contiguous ID vector plus an indptr vector that + * delimits independently sized query rows. Every row is sorted and encoded independently. All + * variable-length portable Roaring streams and their initialized + * `cuco::experimental::roaring_bitmap_ref` objects share one packed device allocation. + * + * For multiple rows, the builder uses indptr directly for segmented device radix sort and + * schedules analysis/encoding over all containers in all rows. A one-row input retains the tuned + * single-allowlist builder. Final encoding and reference initialization remain stream ordered. + * + * IDs must be unique within each row. Setting @p pre_sorted skips sorting and promises that every + * row is strictly increasing; ordering and uniqueness are not checked. Every ID must be smaller + * than dataset_rows. + * + * @see https://github.com/RoaringBitmap/RoaringFormatSpec + * @see + * https://github.com/NVIDIA/cuCollections/blob/6001618aaa7f17ea2bbcd444650e9573c4f3d6c5/include/cuco/roaring_bitmap_ref.cuh + */ +class CUVS_EXPORT roaring_allowlist { + private: + struct impl; + + public: + using key_type = std::uint32_t; + using indptr_type = std::int64_t; + + /** + * @brief Build ragged allowlist rows from contiguous host IDs and row offsets. + * + * `indptr` contains `num_allowlists + 1` entries, starts at zero, is nondecreasing, + * and ends at `ids.extent(0)`. Empty slices are valid allowlists. + */ + static roaring_allowlist from_ids(raft::resources const& res, + std::size_t dataset_rows, + raft::host_vector_view ids, + raft::host_vector_view indptr, + bool pre_sorted = false); + + /** + * @brief Build ragged allowlist rows from contiguous device IDs and row offsets. + * + * The same indptr invariants as the host overload apply. The row offsets are copied to the host + * once for validation, shape-aware dispatch, and exact packed allocation. + * + * The input must remain valid until the construction stream reaches the work enqueued by this + * call. Temporary memory is O(total input IDs + total containers); no dense dataset-sized bitmap + * is materialized. + */ + static roaring_allowlist from_ids( + raft::resources const& res, + std::size_t dataset_rows, + raft::device_vector_view ids, + raft::device_vector_view indptr, + bool pre_sorted = false); + + /** + * @brief Import packed standard 32-bit portable Roaring rows. + * + * `byte_offsets` has `num_allowlists + 1` entries, starts at zero, is nondecreasing, and ends at + * `bytes.extent(0)`. Empty slices represent empty allowlists. Every row is strictly validated on + * the host before its bytes are copied. The host buffers must remain valid until the construction + * stream completes; pinned bytes are recommended when overlap matters. + */ + static roaring_allowlist from_serialized( + raft::resources const& res, + std::size_t dataset_rows, + raft::host_vector_view bytes, + raft::host_vector_view byte_offsets); + + ~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 num_allowlists() const noexcept; + [[nodiscard]] std::size_t dataset_rows() const noexcept; + [[nodiscard]] std::size_t cardinality(std::size_t allowlist_id) const; + [[nodiscard]] bool empty(std::size_t allowlist_id) const; + [[nodiscard]] std::size_t total_cardinality() const noexcept; + + /** @brief Total device bytes retained by packed rows, references, and row pointer metadata. */ + [[nodiscard]] std::size_t size_bytes() const noexcept; + + /** @brief Return a zero-copy view of one row. */ + [[nodiscard]] roaring_allowlist_view view(std::size_t allowlist_id) const; + + /** + * @brief Test a matrix of row IDs and synchronize the resource stream. + * + * `row_ids[q][i]` is tested against allowlist row `q`. Input and output shapes must match, and + * their first extent must equal @ref num_allowlists. + */ + void contains(raft::resources const& res, + raft::device_matrix_view row_ids, + raft::device_matrix_view output) const; + + /** @brief Stream-ordered asynchronous version of @ref contains. */ + void contains_async( + raft::resources const& res, + raft::device_matrix_view row_ids, + raft::device_matrix_view output) const; + + 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 935938c9b0..f2d7f2a5f5 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,91 @@ struct bloom_filter : public base_filter { FilterType get_filter_type() const override { return FilterType::Bloom; } }; +/** + * @brief Reusable per-query mapping to an immutable batch of exact Roaring allowlists. + * + * Entry @c q selects row @c q of the owner. CAGRA retains candidate dataset row @c r when that + * allowlist contains @c r. Constructing from a @c cuvs::core::roaring_allowlist copies only its + * 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, parsing, + * initialization, synchronization, or per-query preprocessing. + * + * @code{.cpp} + * // Flat IDs plus num_queries + 1 row offsets. + * auto allowlists = cuvs::core::roaring_allowlist::from_ids( + * res, dataset_rows, + * raft::make_host_vector_view(allowed_ids.data(), + * allowed_ids.size()), + * raft::make_host_vector_view(indptr.data(), + * indptr.size())); + * std::vector views; + * for (std::size_t q = 0; q < allowlists.num_allowlists(); ++q) { + * views.push_back(allowlists.view(q)); + * } + * auto filter = cuvs::neighbors::filtering::roaring_filter(res, views); + * @endcode + * + * The span overload remains useful when queries reuse rows from several owners or when one query's + * mapping must be replaced without rebuilding encoded allowlists. This filter owns its mapping + * tables and device payload, but not the referenced owner(s), which must outlive the filter and all + * searches using it. Copies are cheap shared handles required by CAGRA query-offset wrappers. + * + * @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 Maximum rejected fraction among all query allowlists. + * + * CAGRA uses this precomputed value when `search_params::filtering_rate` is unset. + */ + [[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/src/core/roaring_allowlist.cu b/cpp/src/core/roaring_allowlist.cu new file mode 100644 index 0000000000..42a2778020 --- /dev/null +++ b/cpp/src/core/roaring_allowlist.cu @@ -0,0 +1,2249 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * 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 + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::core { +namespace { + +/** + * Portable Roaring serialization used by each allowlist + * ===================================================== + * + * This file writes and validates the standard 32-bit portable Roaring format. + * The authoritative format description is: + * + * https://github.com/RoaringBitmap/RoaringFormatSpec + * + * The resulting bytes are consumed on the device by cuCollections: + * + * https://github.com/NVIDIA/cuCollections/blob/6001618aaa7f17ea2bbcd444650e9573c4f3d6c5/include/cuco/roaring_bitmap_ref.cuh + * https://github.com/NVIDIA/cuCollections/blob/6001618aaa7f17ea2bbcd444650e9573c4f3d6c5/include/cuco/detail/roaring_bitmap/util.cuh + * + * All integers below are little-endian. A 32-bit ID is divided into a container + * key and a value: + * + * ID = (uint32_t(key) << 16) | value + * high 16 bits low 16 bits + * + * IDs with the same key belong to one container. Container keys and array + * values are strictly increasing. The portable stream for one nonempty + * allowlist has one of these two headers: + * + * With no run containers, the row is laid out as: + * + * @code{.unparsed} + * uint32 cookie = 12346 + * uint32 N + * descriptor[N] + * uint32 container_offset[N] + * container payloads + * @endcode + * + * With at least one run container, the row is laid out as: + * + * @code{.unparsed} + * uint32 cookie = 12347 | ((N - 1) << 16) + * uint8 run_container_bitmap[ceil(N / 8)] + * descriptor[N] + * uint32 container_offset[N] // present only when N >= 4 + * container payloads + * @endcode + * + * Each four-byte descriptor is: + * + * uint16 key + * uint16 cardinality_minus_one + * + * An offset is measured from the first byte of this stream. Without run + * containers, cardinality selects the payload representation: at most 4096 + * values use an array; more than 4096 use a bitmap. The run-container bitmap + * overrides that choice for marked containers. Its bit order is + * least-significant bit first. + * + * @code{.unparsed} + * array: + * uint16 value[cardinality] + * + * bitmap: + * uint64 words[1024] // 8192 bytes; v is bit (v % 64) of word (v / 64) + * + * run: + * uint16 number_of_runs + * { uint16 start; uint16 length_minus_one; } runs[number_of_runs] + * @endcode + * + * The portable format above describes exactly one bitmap. Query-to-allowlist + * association is a separate concern: `filtering::roaring_filter` stores device + * pointers to already initialized allowlist references. Consequently neither + * the serialized payload nor this metadata is copied or parsed by CAGRA search. + */ + +// Names and thresholds used by the portable format specification. +constexpr std::uint32_t kCookieNoRun = 12346; +constexpr std::uint32_t kCookieRun = 12347; +constexpr std::size_t kArrayCardinality = 4096; +constexpr std::size_t kBitmapBytes = 8192; +constexpr std::size_t kOffsetThreshold = 4; + +using ref_type = cuco::experimental::roaring_bitmap_ref; + +struct row_metadata { + std::size_t cardinality{}; + std::uint32_t max_id{}; + bool empty{true}; +}; + +std::uint16_t read_u16(std::byte const* data, std::size_t size, std::size_t offset) +{ + RAFT_EXPECTS(offset <= size && size - offset >= 2, + "Malformed portable Roaring bitmap: truncated uint16 value."); + return static_cast(std::to_integer(data[offset])) | + static_cast(std::to_integer(data[offset + 1])) << 8; +} + +std::uint32_t read_u32(std::byte const* data, std::size_t size, std::size_t offset) +{ + RAFT_EXPECTS(offset <= size && size - offset >= 4, + "Malformed portable Roaring bitmap: truncated uint32 value."); + std::uint32_t value{}; + for (int i = 0; i < 4; ++i) { + value |= static_cast(std::to_integer(data[offset + i])) << (8 * i); + } + return value; +} + +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."); +} + +enum class container_kind : std::uint8_t { array, bitmap, run }; + +std::size_t align_up(std::size_t offset, std::size_t alignment) +{ + return (offset + alignment - 1) / alignment * alignment; +} + +struct device_build_summary { + std::int64_t cardinality{}; + std::uint64_t payload_bytes{}; + std::uint32_t num_containers{}; + std::uint32_t has_run{}; + std::uint32_t invalid{}; +}; + +struct device_build_result { + rmm::device_uvector storage; + std::size_t serialized_bytes{}; + std::size_t cardinality{}; + bool reference_initialized{}; +}; + +std::size_t reference_offset(std::size_t serialized_bytes) +{ + return align_up(serialized_bytes, alignof(ref_type)); +} + +std::size_t owned_storage_bytes(std::size_t serialized_bytes) +{ + return serialized_bytes == 0 ? 0 : reference_offset(serialized_bytes) + sizeof(ref_type); +} + +constexpr int kBuilderBlockSize = 256; + +// Small allowlists do not benefit from the general builder's device-wide sort, +// two scans, and separate per-stage allocations. At this cardinality every +// portable container is necessarily an array or a run (a bitmap requires more +// than 4096 values in one high-16-bit partition), so one CTA can sort, analyze, +// and later encode the complete row. For a single pre-sorted row the cutoff is lower because the +// general path already avoids its most expensive stage, the device-wide sort. Batched rows use the +// 128-ID capacity because the launch is amortized across the matrix; keep both rules tied to the +// construction benchmark. +constexpr int kSparseBuilderBlockSize = 128; +constexpr std::size_t kSparseBuilderMaxIds = 128; +constexpr std::size_t kSparseBuilderMaxPreSortedIds = 64; +constexpr int kSparseItemsPerThread = + static_cast(kSparseBuilderMaxIds) / kSparseBuilderBlockSize; +static_assert(kSparseBuilderMaxIds % kSparseBuilderBlockSize == 0); + +struct sparse_container_metadata { + std::uint32_t begin{}; + std::uint32_t payload_offset{}; + std::uint16_t runs{}; + container_kind kind{}; + std::uint8_t padding{}; +}; + +static_assert(sizeof(sparse_container_metadata) == 12); + +struct sparse_scratch_layout { + explicit sparse_scratch_layout(std::size_t ids, std::size_t containers, bool store_sorted_ids) + { + sorted_ids_offset = align_up(sizeof(device_build_summary), alignof(std::uint32_t)); + auto offset = sorted_ids_offset + (store_sorted_ids ? ids * sizeof(std::uint32_t) : 0); + metadata_offset = align_up(offset, alignof(sparse_container_metadata)); + bytes = metadata_offset + containers * sizeof(sparse_container_metadata); + } + + std::size_t sorted_ids_offset{}; + std::size_t metadata_offset{}; + std::size_t bytes{}; +}; + +/** + * One allocation for all general-builder temporaries. + * + * The allocation is cardinality/container scaled. The largest CUB workspace is reused by sort, + * boundary selection, and payload scan because those stages are stream ordered. + */ +struct general_scratch_layout { + general_scratch_layout(std::size_t ids, + std::size_t containers, + bool store_sorted_ids, + std::size_t workspace_bytes) + { + std::size_t cursor{}; + auto reserve = [&](std::size_t count, std::size_t item_size, std::size_t alignment) { + auto const result = align_up(cursor, alignment); + cursor = result + count * item_size; + return result; + }; + + if (store_sorted_ids) { + sorted_ids_offset = reserve(ids, sizeof(std::uint32_t), alignof(std::uint32_t)); + } + id_count_offset = reserve(1, sizeof(std::int64_t), alignof(std::int64_t)); + valid_count_offset = reserve(1, sizeof(std::int64_t), alignof(std::int64_t)); + selected_count_offset = reserve(1, sizeof(std::int64_t), alignof(std::int64_t)); + container_starts_offset = reserve(containers, sizeof(std::int64_t), alignof(std::int64_t)); + num_containers_offset = reserve(1, sizeof(std::uint32_t), alignof(std::uint32_t)); + kinds_offset = reserve(containers, sizeof(container_kind), alignof(container_kind)); + payload_sizes_offset = reserve(containers, sizeof(std::uint64_t), alignof(std::uint64_t)); + payload_offsets_offset = reserve(containers, sizeof(std::uint64_t), alignof(std::uint64_t)); + has_run_offset = reserve(1, sizeof(std::uint32_t), alignof(std::uint32_t)); + summary_offset = reserve(1, sizeof(device_build_summary), alignof(device_build_summary)); + workspace_offset = reserve(workspace_bytes, sizeof(cuda::std::byte), alignof(std::max_align_t)); + bytes = cursor; + } + + std::size_t sorted_ids_offset{}; + std::size_t id_count_offset{}; + std::size_t valid_count_offset{}; + std::size_t selected_count_offset{}; + std::size_t container_starts_offset{}; + std::size_t num_containers_offset{}; + std::size_t kinds_offset{}; + std::size_t payload_sizes_offset{}; + std::size_t payload_offsets_offset{}; + std::size_t has_run_offset{}; + std::size_t summary_offset{}; + std::size_t workspace_offset{}; + std::size_t bytes{}; +}; + +int grid_size_for(std::size_t count) +{ + auto const blocks = (count + kBuilderBlockSize - 1) / kBuilderBlockSize; + return static_cast(std::min(blocks, 65535)); +} + +__device__ void write_u16(cuda::std::byte* output, std::size_t offset, std::uint16_t value) +{ + auto* bytes = reinterpret_cast(output); + bytes[offset] = static_cast(value); + bytes[offset + 1] = static_cast(value >> 8); +} + +__device__ void write_u32(cuda::std::byte* output, std::size_t offset, std::uint32_t value) +{ + auto* bytes = reinterpret_cast(output); + for (int byte = 0; byte < 4; ++byte) { + bytes[offset + byte] = static_cast(value >> (8 * byte)); + } +} + +__device__ void write_u64(cuda::std::byte* output, std::size_t offset, std::uint64_t value) +{ + auto* bytes = reinterpret_cast(output); + for (int byte = 0; byte < 8; ++byte) { + bytes[offset + byte] = static_cast(value >> (8 * byte)); + } +} + +/** Find the sorted prefix that lies inside the logical dataset shape. */ +__global__ void find_valid_count_kernel(std::uint32_t const* ids, + std::int64_t const* id_count, + std::uint64_t dataset_rows, + std::int64_t* valid_count) +{ + if (blockIdx.x != 0 || threadIdx.x != 0) { return; } + std::int64_t first{}; + auto last = *id_count; + while (first < last) { + auto const middle = first + (last - first) / 2; + if (static_cast(ids[middle]) < dataset_rows) { + first = middle + 1; + } else { + last = middle; + } + } + *valid_count = first; +} + +/** Select the first sorted ID belonging to every high-16-bit container. */ +struct is_container_start { + std::uint32_t const* ids{}; + std::int64_t const* valid_count{}; + + __device__ bool operator()(std::int64_t i) const + { + auto const count = *valid_count; + return i < count && (i == 0 || (ids[i - 1] >> 16) != (ids[i] >> 16)); + } +}; + +__global__ void narrow_container_count_kernel(std::int64_t const* selected_count, + std::uint32_t* num_containers) +{ + if (blockIdx.x == 0 && threadIdx.x == 0) { + *num_containers = static_cast(*selected_count); + } +} + +/** + * Count consecutive runs and select the smallest legal portable payload for + * each container. + * + * Each block owns one container. Threads independently identify run starts in + * the sorted slice, then a block reduction produces the exact run count. This + * uses O(number of input IDs) scratch for sorting and scans; it never + * constructs a dense dataset-sized bitmap. + */ +__global__ void analyze_containers_kernel(std::uint32_t const* ids, + std::int64_t const* id_count, + std::int64_t const* container_starts, + std::uint32_t const* num_containers, + container_kind* kinds, + std::uint64_t* payload_sizes, + std::uint32_t* has_run) +{ + auto const container = static_cast(blockIdx.x); + auto const count = *num_containers; + if (container >= count) { return; } + + auto const begin = container_starts[container]; + auto const end = container + 1 < count ? container_starts[container + 1] : *id_count; + std::uint32_t local_runs{}; + for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { + local_runs += + i == begin || static_cast(ids[i]) != static_cast(ids[i - 1]) + 1 + ? 1u + : 0u; + } + + using block_reduce = cub::BlockReduce; + __shared__ typename block_reduce::TempStorage reduction_storage; + auto const runs = block_reduce(reduction_storage).Sum(local_runs); + if (threadIdx.x != 0) { return; } + + auto const cardinality = static_cast(end - begin); + auto const normal_size = + cardinality <= kArrayCardinality ? cardinality * sizeof(std::uint16_t) : kBitmapBytes; + auto const run_size = sizeof(std::uint16_t) + runs * 2 * sizeof(std::uint16_t); + if (run_size < normal_size) { + kinds[container] = container_kind::run; + payload_sizes[container] = run_size; + atomicExch(has_run, 1u); + } else if (cardinality <= kArrayCardinality) { + kinds[container] = container_kind::array; + payload_sizes[container] = normal_size; + } else { + kinds[container] = container_kind::bitmap; + payload_sizes[container] = normal_size; + } +} + +/** Collect the scalar results needed to allocate the exact final portable byte + * stream. */ +__global__ void finish_device_analysis_kernel(std::int64_t const* id_count, + std::int64_t const* valid_count, + std::uint32_t const* num_containers, + std::uint32_t const* has_run, + std::uint64_t const* payload_sizes, + std::uint64_t const* payload_offsets, + device_build_summary* summary) +{ + if (blockIdx.x != 0 || threadIdx.x != 0) { return; } + auto const cardinality = *id_count; + auto const containers = *num_containers; + summary->cardinality = cardinality; + summary->num_containers = containers; + summary->has_run = *has_run; + summary->payload_bytes = + containers == 0 ? 0 : payload_offsets[containers - 1] + payload_sizes[containers - 1]; + summary->invalid = cardinality != *valid_count; +} + +__host__ __device__ std::size_t portable_header_size(std::uint32_t num_containers, bool has_run) +{ + if (!has_run) { + return 2 * sizeof(std::uint32_t) + + num_containers * (2 * sizeof(std::uint16_t) + sizeof(std::uint32_t)); + } + auto const run_bitmap_bytes = (num_containers + 7) / 8; + return sizeof(std::uint32_t) + run_bitmap_bytes + num_containers * 2 * sizeof(std::uint16_t) + + (num_containers >= kOffsetThreshold ? num_containers * sizeof(std::uint32_t) : 0); +} + +/** Write the cookie, run bitmap, descriptors, and portable container-offset + * table. */ +__global__ void encode_header_kernel(std::uint32_t const* ids, + std::int64_t const* id_count, + std::int64_t const* container_starts, + std::uint32_t num_containers, + container_kind const* kinds, + std::uint64_t const* payload_offsets, + std::size_t header_size, + bool has_run, + cuda::std::byte* output) +{ + auto const thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + auto const stride = static_cast(gridDim.x) * blockDim.x; + auto const run_bitmap_bytes = has_run ? (num_containers + 7) / 8 : 0; + auto const descriptor_offset = + has_run ? sizeof(std::uint32_t) + run_bitmap_bytes : 2 * sizeof(std::uint32_t); + auto const offsets_offset = descriptor_offset + num_containers * 2 * sizeof(std::uint16_t); + bool const store_offsets = !has_run || num_containers >= kOffsetThreshold; + + if (thread == 0) { + if (has_run) { + write_u32(output, 0, kCookieRun | ((num_containers - 1) << 16)); + } else { + write_u32(output, 0, kCookieNoRun); + write_u32(output, sizeof(std::uint32_t), num_containers); + } + } + + // Write every run-bitmap byte directly. This initializes unused high bits to zero and avoids a + // memset of the complete serialized allocation. + auto* output_bytes = reinterpret_cast(output); + for (auto byte = thread; byte < run_bitmap_bytes; byte += stride) { + std::uint8_t value{}; + for (std::uint32_t bit = 0; bit < 8; ++bit) { + auto const container = static_cast(byte * 8 + bit); + if (container < num_containers && kinds[container] == container_kind::run) { + value |= static_cast(1u << bit); + } + } + output_bytes[sizeof(std::uint32_t) + byte] = value; + } + + for (auto container = thread; container < num_containers; container += stride) { + auto const begin = container_starts[container]; + auto const end = container + 1 < num_containers ? container_starts[container + 1] : *id_count; + auto const descriptor = descriptor_offset + container * 2 * sizeof(std::uint16_t); + write_u16(output, descriptor, static_cast(ids[begin] >> 16)); + write_u16( + output, descriptor + sizeof(std::uint16_t), static_cast(end - begin - 1)); + if (store_offsets) { + write_u32(output, + offsets_offset + container * sizeof(std::uint32_t), + static_cast(header_size + payload_offsets[container])); + } + } +} + +/** Encode array, bitmap, and run payloads directly into their final device offsets. */ +__global__ void encode_payloads_kernel(std::uint32_t const* ids, + std::int64_t const* id_count, + std::int64_t const* container_starts, + std::uint32_t num_containers, + container_kind const* kinds, + std::uint64_t const* payload_sizes, + std::uint64_t const* payload_offsets, + std::size_t header_size, + cuda::std::byte* output) +{ + auto const container = static_cast(blockIdx.x); + if (container >= num_containers) { return; } + auto const begin = container_starts[container]; + auto const end = container + 1 < num_containers ? container_starts[container + 1] : *id_count; + auto const payload = header_size + payload_offsets[container]; + + if (kinds[container] == container_kind::array) { + for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { + write_u16(output, + payload + static_cast(i - begin) * sizeof(std::uint16_t), + static_cast(ids[i] & 0xffffu)); + } + return; + } + + using run_scan = cub::BlockScan; + union payload_scratch { + std::uint64_t bitmap_words[kBitmapBytes / sizeof(std::uint64_t)]; + typename run_scan::TempStorage run_scan_storage; + }; + __shared__ payload_scratch scratch; + __shared__ std::uint32_t run_base; + __shared__ std::uint32_t tile_runs; + + if (kinds[container] == container_kind::bitmap) { + constexpr std::uint32_t words = kBitmapBytes / sizeof(std::uint64_t); + for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { + scratch.bitmap_words[word] = 0; + } + __syncthreads(); + for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { + auto const lower = ids[i] & 0xffffu; + atomicOr(reinterpret_cast(&scratch.bitmap_words[lower / 64]), + static_cast(std::uint64_t{1} << (lower % 64))); + } + __syncthreads(); + for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { + write_u64(output, payload + word * sizeof(std::uint64_t), scratch.bitmap_words[word]); + } + return; + } + + auto const num_runs = static_cast( + (payload_sizes[container] - sizeof(std::uint16_t)) / (2 * sizeof(std::uint16_t))); + if (threadIdx.x == 0) { + write_u16(output, payload, num_runs); + run_base = 0; + } + __syncthreads(); + + // A block scan assigns stable output positions to run starts in each tile. Each run-start thread + // walks only its own run, so total work remains O(container cardinality) while high-run-count + // containers use the complete CTA instead of one serial thread. + for (auto tile = begin; tile < end; tile += blockDim.x) { + auto const i = tile + threadIdx.x; + std::uint32_t const is_run_start = + i < end && (i == begin || + static_cast(ids[i]) != static_cast(ids[i - 1]) + 1) + ? 1u + : 0u; + std::uint32_t run_rank{}; + std::uint32_t block_runs{}; + run_scan(scratch.run_scan_storage).ExclusiveSum(is_run_start, run_rank, block_runs); + if (threadIdx.x == 0) { tile_runs = block_runs; } + __syncthreads(); + + if (is_run_start != 0) { + auto j = i + 1; + while (j < end && + static_cast(ids[j]) == static_cast(ids[j - 1]) + 1) { + ++j; + } + auto const start = static_cast(ids[i] & 0xffffu); + auto const last = static_cast(ids[j - 1] & 0xffffu); + auto const run_offset = + payload + sizeof(std::uint16_t) + + static_cast(run_base + run_rank) * 2 * sizeof(std::uint16_t); + write_u16(output, run_offset, start); + write_u16( + output, run_offset + sizeof(std::uint16_t), static_cast(last - start)); + } + __syncthreads(); + if (threadIdx.x == 0) { run_base += tile_runs; } + __syncthreads(); + } +} + +/** + * Sort and analyze an entire sparse allowlist in one CTA. + * + * The unsorted specialization uses a blocked 128 x 1 radix sort. Padding uses UINT32_MAX; + * writing only the first `size` sorted items is still correct when UINT32_MAX itself is a valid ID + * because all padding values compare equal to that final real value. Thread zero then walks at most + * 128 normalized IDs to + * build compact per-container metadata and the exact serialized payload size. + * + * Inputs are promised unique by the public API. This kernel deliberately does + * not spend work or storage checking or collapsing duplicates. + */ +template +__global__ void analyze_sparse_ids_kernel(std::uint32_t const* ids, + std::uint32_t size, + std::uint64_t dataset_rows, + std::uint32_t* sorted_ids, + sparse_container_metadata* metadata, + device_build_summary* summary) +{ + if constexpr (SortInput) { + using block_sort = + cub::BlockRadixSort; + __shared__ typename block_sort::TempStorage sort_storage; + std::uint32_t thread_ids[kSparseItemsPerThread]; + +#pragma unroll + for (int item = 0; item < kSparseItemsPerThread; ++item) { + auto const index = static_cast(threadIdx.x) * kSparseItemsPerThread + item; + thread_ids[item] = index < size ? ids[index] : std::numeric_limits::max(); + } + block_sort(sort_storage).Sort(thread_ids); +#pragma unroll + for (int item = 0; item < kSparseItemsPerThread; ++item) { + auto const index = static_cast(threadIdx.x) * kSparseItemsPerThread + item; + if (index < size) { sorted_ids[index] = thread_ids[item]; } + } + __syncthreads(); + } + + if (threadIdx.x != 0) { return; } + auto const* normalized_ids = SortInput ? sorted_ids : ids; + + summary->cardinality = size; + summary->payload_bytes = 0; + summary->num_containers = 0; + summary->has_run = 0; + summary->invalid = 0; + + // Validate before writing container metadata. For a valid row, the logical + // dataset shape bounds the number of containers allocated by the host. + for (std::uint32_t i = 0; i < size; ++i) { + if (static_cast(normalized_ids[i]) >= dataset_rows) { + summary->invalid = 1; + return; + } + } + + std::uint32_t container{}; + std::uint32_t begin{}; + std::uint32_t payload_offset{}; + while (begin < size) { + auto const key = normalized_ids[begin] >> 16; + auto end = begin + 1; + std::uint32_t runs{1}; + while (end < size && (normalized_ids[end] >> 16) == key) { + runs += static_cast(normalized_ids[end]) != + static_cast(normalized_ids[end - 1]) + 1 + ? 1u + : 0u; + ++end; + } + + auto const cardinality = end - begin; + auto const array_size = cardinality * sizeof(std::uint16_t); + auto const run_size = sizeof(std::uint16_t) + runs * 2 * sizeof(std::uint16_t); + auto const use_run = run_size < array_size; + metadata[container] = + sparse_container_metadata{begin, + payload_offset, + static_cast(runs), + use_run ? container_kind::run : container_kind::array, + 0}; + payload_offset += use_run ? run_size : array_size; + summary->has_run |= use_run ? 1u : 0u; + ++container; + begin = end; + } + + summary->payload_bytes = payload_offset; + summary->num_containers = container; +} + +/** + * Encode a sparse row in one CTA after the host has allocated the exact byte + * count reported by `analyze_sparse_ids_kernel`. + * + * Thread zero writes every header byte, including the run bitmap, so this path + * needs no output memset. Array values are striped across the CTA; thread zero + * writes the comparatively small run payloads. Bitmap payloads cannot occur + * below the sparse cardinality threshold. + */ +__global__ void encode_sparse_row_kernel(std::uint32_t const* ids, + std::uint32_t cardinality, + sparse_container_metadata const* metadata, + std::uint32_t num_containers, + std::size_t header_size, + bool has_run, + cuda::std::byte* output, + ref_type* reference) +{ + bool const store_offsets = !has_run || num_containers >= kOffsetThreshold; + std::size_t descriptor_offset{}; + std::size_t offsets_offset{}; + + if (threadIdx.x == 0) { + if (has_run) { + write_u32(output, 0, kCookieRun | ((num_containers - 1) << 16)); + auto const run_bitmap_bytes = (num_containers + 7) / 8; + for (std::uint32_t byte = 0; byte < run_bitmap_bytes; ++byte) { + std::uint8_t value{}; + for (std::uint32_t bit = 0; bit < 8; ++bit) { + auto const container = byte * 8 + bit; + if (container < num_containers && metadata[container].kind == container_kind::run) { + value |= static_cast(1u << bit); + } + } + reinterpret_cast(output)[sizeof(std::uint32_t) + byte] = value; + } + descriptor_offset = sizeof(std::uint32_t) + run_bitmap_bytes; + offsets_offset = descriptor_offset + num_containers * 2 * sizeof(std::uint16_t); + } else { + write_u32(output, 0, kCookieNoRun); + write_u32(output, sizeof(std::uint32_t), num_containers); + descriptor_offset = 2 * sizeof(std::uint32_t); + offsets_offset = descriptor_offset + num_containers * 2 * sizeof(std::uint16_t); + } + + for (std::uint32_t container = 0; container < num_containers; ++container) { + auto const begin = metadata[container].begin; + auto const end = container + 1 < num_containers ? metadata[container + 1].begin : cardinality; + auto const descriptor = descriptor_offset + container * 2 * sizeof(std::uint16_t); + write_u16(output, descriptor, static_cast(ids[begin] >> 16)); + write_u16( + output, descriptor + sizeof(std::uint16_t), static_cast(end - begin - 1)); + if (store_offsets) { + write_u32(output, + offsets_offset + container * sizeof(std::uint32_t), + static_cast(header_size + metadata[container].payload_offset)); + } + } + } + + for (std::uint32_t container = 0; container < num_containers; ++container) { + auto const begin = metadata[container].begin; + auto const end = container + 1 < num_containers ? metadata[container + 1].begin : cardinality; + auto const payload = header_size + metadata[container].payload_offset; + if (metadata[container].kind == container_kind::array) { + for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { + write_u16(output, + payload + static_cast(i - begin) * sizeof(std::uint16_t), + static_cast(ids[i] & 0xffffu)); + } + continue; + } + + if (threadIdx.x == 0) { + write_u16(output, payload, metadata[container].runs); + std::uint16_t run_index{}; + for (auto i = begin; i < end;) { + auto const start = static_cast(ids[i] & 0xffffu); + auto j = i + 1; + while (j < end && + static_cast(ids[j]) == static_cast(ids[j - 1]) + 1) { + ++j; + } + auto const last = static_cast(ids[j - 1] & 0xffffu); + auto const run_offset = payload + sizeof(std::uint16_t) + + static_cast(run_index) * 2 * sizeof(std::uint16_t); + write_u16(output, run_offset, start); + write_u16( + output, run_offset + sizeof(std::uint16_t), static_cast(last - start)); + ++run_index; + i = j; + } + } + } + __syncthreads(); + if (threadIdx.x == 0) { ::new (static_cast(reference)) ref_type{output}; } +} + +device_build_result build_sparse_from_device_ids( + raft::resources const& res, + std::size_t dataset_rows, + raft::device_vector_view ids, + bool pre_sorted) +{ + common::nvtx::range build_scope("roaring_allowlist::build_sparse"); + auto const stream = raft::resource::get_cuda_stream(res); + auto const size = static_cast(ids.extent(0)); + auto const num_chunks = + (static_cast(dataset_rows) + (std::uint64_t{1} << 16) - 1) >> 16; + auto const max_containers = std::min(size, static_cast(num_chunks)); + sparse_scratch_layout const layout{size, max_containers, !pre_sorted}; + rmm::device_uvector scratch(layout.bytes, stream); + + auto* summary = reinterpret_cast(scratch.data()); + auto* sorted_ids = + pre_sorted ? nullptr + : reinterpret_cast(scratch.data() + layout.sorted_ids_offset); + auto* metadata = + reinterpret_cast(scratch.data() + layout.metadata_offset); + + { + common::nvtx::range stage_scope( + "roaring_allowlist::sparse_analysis"); + if (pre_sorted) { + analyze_sparse_ids_kernel + <<<1, kSparseBuilderBlockSize, 0, stream>>>(ids.data_handle(), + static_cast(size), + dataset_rows, + nullptr, + metadata, + summary); + } else { + analyze_sparse_ids_kernel + <<<1, kSparseBuilderBlockSize, 0, stream>>>(ids.data_handle(), + static_cast(size), + dataset_rows, + sorted_ids, + metadata, + summary); + } + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + device_build_summary host_summary; + { + common::nvtx::range stage_scope("roaring_allowlist::size_readback"); + RAFT_CUDA_TRY(cudaMemcpyAsync( + &host_summary, summary, sizeof(host_summary), cudaMemcpyDeviceToHost, stream)); + raft::resource::sync_stream(res); + } + RAFT_EXPECTS(host_summary.invalid == 0, + "Roaring allowlist ID must be smaller than dataset_rows."); + RAFT_EXPECTS(host_summary.cardinality > 0 && host_summary.num_containers > 0, + "Internal error: nonempty sparse Roaring input produced an empty device build."); + + auto const header_size = + portable_header_size(host_summary.num_containers, host_summary.has_run != 0); + auto const serialized_bytes = header_size + static_cast(host_summary.payload_bytes); + rmm::device_uvector output(0, stream); + { + common::nvtx::range stage_scope( + "roaring_allowlist::final_allocation"); + output.resize(owned_storage_bytes(serialized_bytes), stream); + } + auto const* normalized_ids = pre_sorted ? ids.data_handle() : sorted_ids; + common::nvtx::range encode_scope( + "roaring_allowlist::sparse_encode_and_ref"); + encode_sparse_row_kernel<<<1, kSparseBuilderBlockSize, 0, stream>>>( + normalized_ids, + static_cast(size), + metadata, + host_summary.num_containers, + header_size, + host_summary.has_run != 0, + output.data(), + reinterpret_cast(output.data() + reference_offset(serialized_bytes))); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + return {std::move(output), serialized_bytes, size, true}; +} + +/** + * Build a standard portable Roaring row from device IDs. + * + * Very sparse rows use the single-CTA builder above. Larger rows use one ID + * array for device-wide radix sorting. Both pre-sorted paths + * use the caller's strictly increasing IDs directly and allocate no normalization ID array. + * Boundary scan data, CUB workspace, and O(min(input IDs, 65536)) container metadata remain + * cardinality-scaled. In particular, the builder never allocates storage proportional to + * `dataset_rows` bits. + * + * Pre-sorted ordering and uniqueness are unchecked caller promises. + */ +device_build_result build_from_device_ids( + raft::resources const& res, + std::size_t dataset_rows, + 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 {rmm::device_uvector(0, stream), 0, 0, false}; } + auto const sparse_cutoff = pre_sorted ? kSparseBuilderMaxPreSortedIds : kSparseBuilderMaxIds; + if (size <= sparse_cutoff) { + return build_sparse_from_device_ids(res, dataset_rows, ids, pre_sorted); + } + + auto const num_chunks = + (static_cast(dataset_rows) + (std::uint64_t{1} << 16) - 1) >> 16; + auto const max_containers = std::min(size, static_cast(num_chunks)); + auto const item_count = static_cast(size); + auto const container_slots = static_cast(max_containers); + constexpr int sort_end_bit = std::numeric_limits::digits; + + std::size_t sort_workspace_bytes{}; + std::size_t select_workspace_bytes{}; + std::size_t payload_scan_workspace_bytes{}; + if (!pre_sorted) { + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortKeys(nullptr, + sort_workspace_bytes, + ids.data_handle(), + static_cast(nullptr), + item_count, + 0, + sort_end_bit, + stream)); + } + auto counting = thrust::make_counting_iterator(0); + RAFT_CUDA_TRY(cub::DeviceSelect::If(nullptr, + select_workspace_bytes, + counting, + static_cast(nullptr), + static_cast(nullptr), + item_count, + is_container_start{ids.data_handle(), nullptr}, + stream)); + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, + payload_scan_workspace_bytes, + static_cast(nullptr), + static_cast(nullptr), + container_slots, + stream)); + auto const workspace_bytes = + std::max(sort_workspace_bytes, std::max(select_workspace_bytes, payload_scan_workspace_bytes)); + general_scratch_layout const layout{size, max_containers, !pre_sorted, workspace_bytes}; + rmm::device_uvector scratch(layout.bytes, stream); + + auto* sorted_ids = + pre_sorted ? nullptr + : reinterpret_cast(scratch.data() + layout.sorted_ids_offset); + auto* id_count = reinterpret_cast(scratch.data() + layout.id_count_offset); + auto* valid_count = reinterpret_cast(scratch.data() + layout.valid_count_offset); + auto* selected_count = + reinterpret_cast(scratch.data() + layout.selected_count_offset); + auto* container_starts = + reinterpret_cast(scratch.data() + layout.container_starts_offset); + auto* num_containers = + reinterpret_cast(scratch.data() + layout.num_containers_offset); + auto* kinds = reinterpret_cast(scratch.data() + layout.kinds_offset); + auto* payload_sizes = + reinterpret_cast(scratch.data() + layout.payload_sizes_offset); + auto* payload_offsets = + reinterpret_cast(scratch.data() + layout.payload_offsets_offset); + auto* has_run = reinterpret_cast(scratch.data() + layout.has_run_offset); + auto* device_summary = + reinterpret_cast(scratch.data() + layout.summary_offset); + auto* workspace = scratch.data() + layout.workspace_offset; + + if (!pre_sorted) { + common::nvtx::range stage_scope("roaring_allowlist::radix_sort"); + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortKeys(workspace, + sort_workspace_bytes, + ids.data_handle(), + sorted_ids, + item_count, + 0, + sort_end_bit, + stream)); + } + auto const* normalized_ids = pre_sorted ? ids.data_handle() : sorted_ids; + RAFT_CUDA_TRY( + cudaMemcpyAsync(id_count, &item_count, sizeof(item_count), cudaMemcpyHostToDevice, stream)); + + { + common::nvtx::range stage_scope( + "roaring_allowlist::container_discovery"); + find_valid_count_kernel<<<1, 1, 0, stream>>>( + normalized_ids, id_count, dataset_rows, valid_count); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + RAFT_CUDA_TRY(cub::DeviceSelect::If(workspace, + select_workspace_bytes, + counting, + container_starts, + selected_count, + item_count, + is_container_start{normalized_ids, valid_count}, + stream)); + narrow_container_count_kernel<<<1, 1, 0, stream>>>(selected_count, num_containers); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + { + common::nvtx::range stage_scope( + "roaring_allowlist::container_analysis"); + RAFT_CUDA_TRY( + cudaMemsetAsync(payload_sizes, 0, max_containers * sizeof(std::uint64_t), stream)); + RAFT_CUDA_TRY(cudaMemsetAsync(has_run, 0, sizeof(std::uint32_t), stream)); + analyze_containers_kernel<<(max_containers), + kBuilderBlockSize, + 0, + stream>>>( + normalized_ids, valid_count, container_starts, num_containers, kinds, payload_sizes, has_run); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(workspace, + payload_scan_workspace_bytes, + payload_sizes, + payload_offsets, + container_slots, + stream)); + finish_device_analysis_kernel<<<1, 1, 0, stream>>>(id_count, + valid_count, + num_containers, + has_run, + payload_sizes, + payload_offsets, + device_summary); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + device_build_summary summary; + { + common::nvtx::range stage_scope("roaring_allowlist::size_readback"); + RAFT_CUDA_TRY( + cudaMemcpyAsync(&summary, device_summary, sizeof(summary), cudaMemcpyDeviceToHost, stream)); + raft::resource::sync_stream(res); + } + RAFT_EXPECTS(summary.invalid == 0, "Roaring allowlist ID must be smaller than dataset_rows."); + RAFT_EXPECTS(summary.cardinality > 0 && summary.num_containers > 0, + "Internal error: nonempty Roaring input produced an empty device build."); + + auto const header_size = portable_header_size(summary.num_containers, summary.has_run != 0); + RAFT_EXPECTS(summary.payload_bytes <= std::numeric_limits::max() - header_size, + "Portable Roaring row exceeds the 32-bit offset range."); + auto const serialized_bytes = header_size + static_cast(summary.payload_bytes); + rmm::device_uvector output(0, stream); + { + common::nvtx::range stage_scope( + "roaring_allowlist::final_allocation"); + output.resize(owned_storage_bytes(serialized_bytes), stream); + } + auto const header_items = + std::max(summary.num_containers, (summary.num_containers + 7) / 8); + { + common::nvtx::range stage_scope("roaring_allowlist::header_encode"); + encode_header_kernel<<>>( + normalized_ids, + valid_count, + container_starts, + summary.num_containers, + kinds, + payload_offsets, + header_size, + summary.has_run != 0, + output.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + common::nvtx::range payload_scope( + "roaring_allowlist::payload_encode"); + encode_payloads_kernel<<>>( + normalized_ids, + valid_count, + container_starts, + summary.num_containers, + kinds, + payload_sizes, + payload_offsets, + header_size, + output.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + return { + std::move(output), serialized_bytes, static_cast(summary.cardinality), false}; +} + +struct batched_device_build_result { + rmm::device_uvector storage; + std::vector row_offsets; + std::vector reference_offsets; + std::vector serialized_bytes; + std::vector cardinalities; + bool references_initialized{}; +}; + +struct packed_rows_layout { + std::vector row_offsets; + std::vector reference_offsets; + std::size_t references_offset{}; + std::size_t bytes{}; +}; + +packed_rows_layout make_packed_rows_layout(std::vector const& serialized_bytes) +{ + packed_rows_layout result; + result.row_offsets.resize(serialized_bytes.size()); + result.reference_offsets.resize(serialized_bytes.size()); + std::size_t cursor{}; + bool any_nonempty{}; + for (std::size_t row = 0; row < serialized_bytes.size(); ++row) { + if (serialized_bytes[row] == 0) { continue; } + any_nonempty = true; + cursor = align_up(cursor, alignof(std::max_align_t)); + result.row_offsets[row] = cursor; + cursor += serialized_bytes[row]; + } + if (!any_nonempty) { return result; } + result.references_offset = align_up(cursor, alignof(ref_type)); + for (std::size_t row = 0; row < serialized_bytes.size(); ++row) { + result.reference_offsets[row] = result.references_offset + row * sizeof(ref_type); + } + result.bytes = result.references_offset + serialized_bytes.size() * sizeof(ref_type); + return result; +} + +struct batch_general_scratch_layout { + batch_general_scratch_layout(std::size_t total_ids, + std::size_t rows, + std::size_t max_containers, + bool store_sorted_ids, + std::size_t workspace_bytes) + { + std::size_t cursor{}; + auto reserve = [&](std::size_t count, std::size_t item_size, std::size_t alignment) { + auto const result = align_up(cursor, alignment); + cursor = result + count * item_size; + return result; + }; + if (store_sorted_ids) { + sorted_ids_offset = reserve(total_ids, sizeof(std::uint32_t), alignof(std::uint32_t)); + } + valid_counts_offset = reserve(rows, sizeof(std::int64_t), alignof(std::int64_t)); + flags_offset = reserve(total_ids, sizeof(std::uint8_t), alignof(std::uint8_t)); + selected_count_offset = reserve(1, sizeof(std::int64_t), alignof(std::int64_t)); + container_starts_offset = reserve(max_containers, sizeof(std::int64_t), alignof(std::int64_t)); + row_container_offsets_offset = reserve(rows + 1, sizeof(std::int64_t), alignof(std::int64_t)); + container_rows_offset = reserve(max_containers, sizeof(std::uint32_t), alignof(std::uint32_t)); + kinds_offset = reserve(max_containers, sizeof(container_kind), alignof(container_kind)); + payload_sizes_offset = reserve(max_containers, sizeof(std::uint64_t), alignof(std::uint64_t)); + payload_offsets_offset = reserve(max_containers, sizeof(std::uint64_t), alignof(std::uint64_t)); + has_run_offset = reserve(rows, sizeof(std::uint32_t), alignof(std::uint32_t)); + summaries_offset = reserve(rows, sizeof(device_build_summary), alignof(device_build_summary)); + output_offsets_offset = reserve(rows, sizeof(std::uint64_t), alignof(std::uint64_t)); + workspace_offset = reserve(workspace_bytes, sizeof(cuda::std::byte), alignof(std::max_align_t)); + bytes = cursor; + } + + std::size_t sorted_ids_offset{}; + std::size_t valid_counts_offset{}; + std::size_t flags_offset{}; + std::size_t selected_count_offset{}; + std::size_t container_starts_offset{}; + std::size_t row_container_offsets_offset{}; + std::size_t container_rows_offset{}; + std::size_t kinds_offset{}; + std::size_t payload_sizes_offset{}; + std::size_t payload_offsets_offset{}; + std::size_t has_run_offset{}; + std::size_t summaries_offset{}; + std::size_t output_offsets_offset{}; + std::size_t workspace_offset{}; + std::size_t bytes{}; +}; + +__global__ void find_batch_valid_counts_kernel(std::uint32_t const* ids, + std::int64_t const* indptr, + std::int64_t rows, + std::uint64_t dataset_rows, + std::int64_t* valid_counts) +{ + auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (row >= rows) { return; } + auto const begin = indptr[row]; + auto const width = indptr[row + 1] - begin; + std::int64_t first{}; + auto last = width; + while (first < last) { + auto const middle = first + (last - first) / 2; + if (static_cast(ids[begin + middle]) < dataset_rows) { + first = middle + 1; + } else { + last = middle; + } + } + valid_counts[row] = first; +} + +__global__ void mark_batch_container_starts_kernel(std::uint32_t const* ids, + std::int64_t const* indptr, + std::int64_t const* valid_counts, + std::int64_t rows, + std::uint8_t* flags) +{ + auto const row = static_cast(blockIdx.x); + if (row >= rows) { return; } + auto const begin = indptr[row]; + auto const width = indptr[row + 1] - begin; + auto const valid = valid_counts[row]; + for (std::int64_t column = threadIdx.x; column < width; column += blockDim.x) { + auto const index = begin + column; + flags[index] = + column < valid && (column == 0 || (ids[index - 1] >> 16) != (ids[index] >> 16)) ? 1 : 0; + } +} + +__global__ void find_row_container_offsets_kernel(std::int64_t const* container_starts, + std::int64_t const* selected_count, + std::int64_t const* indptr, + std::int64_t rows, + std::int64_t* row_offsets) +{ + auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (row > rows) { return; } + auto const target = indptr[row]; + std::int64_t first{}; + auto last = *selected_count; + while (first < last) { + auto const middle = first + (last - first) / 2; + if (container_starts[middle] < target) { + first = middle + 1; + } else { + last = middle; + } + } + row_offsets[row] = first; +} + +__global__ void fill_container_rows_kernel(std::int64_t const* row_container_offsets, + std::int64_t rows, + std::uint32_t* container_rows) +{ + auto const row = static_cast(blockIdx.x); + if (row >= rows) { return; } + for (auto container = row_container_offsets[row] + threadIdx.x; + container < row_container_offsets[row + 1]; + container += blockDim.x) { + container_rows[container] = static_cast(row); + } +} + +__global__ void analyze_batch_containers_kernel(std::uint32_t const* ids, + std::int64_t const* indptr, + std::int64_t const* valid_counts, + std::int64_t const* container_starts, + std::int64_t const* selected_count, + std::uint32_t const* container_rows, + container_kind* kinds, + std::uint64_t* payload_sizes, + std::uint32_t* has_run) +{ + auto const container = static_cast(blockIdx.x); + auto const count = *selected_count; + if (container >= count) { return; } + auto const row = static_cast(container_rows[container]); + auto const begin = container_starts[container]; + auto const row_end = indptr[row] + valid_counts[row]; + auto const end = container + 1 < count && container_rows[container + 1] == row + ? container_starts[container + 1] + : row_end; + std::uint32_t local_runs{}; + for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { + local_runs += + i == begin || static_cast(ids[i]) != static_cast(ids[i - 1]) + 1 + ? 1u + : 0u; + } + using block_reduce = cub::BlockReduce; + __shared__ typename block_reduce::TempStorage reduction_storage; + auto const runs = block_reduce(reduction_storage).Sum(local_runs); + if (threadIdx.x != 0) { return; } + auto const cardinality = static_cast(end - begin); + auto const normal_size = + cardinality <= kArrayCardinality ? cardinality * sizeof(std::uint16_t) : kBitmapBytes; + auto const run_size = sizeof(std::uint16_t) + runs * 2 * sizeof(std::uint16_t); + if (run_size < normal_size) { + kinds[container] = container_kind::run; + payload_sizes[container] = run_size; + atomicExch(has_run + row, 1u); + } else if (cardinality <= kArrayCardinality) { + kinds[container] = container_kind::array; + payload_sizes[container] = normal_size; + } else { + kinds[container] = container_kind::bitmap; + payload_sizes[container] = normal_size; + } +} + +__global__ void finish_batch_rows_kernel(std::int64_t rows, + std::int64_t const* indptr, + std::int64_t const* valid_counts, + std::int64_t const* row_container_offsets, + std::uint32_t const* has_run, + std::uint64_t const* payload_sizes, + std::uint64_t const* payload_offsets, + device_build_summary* summaries) +{ + auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (row >= rows) { return; } + auto const begin = row_container_offsets[row]; + auto const end = row_container_offsets[row + 1]; + auto const cardinality = indptr[row + 1] - indptr[row]; + auto& summary = summaries[row]; + summary.cardinality = cardinality; + summary.num_containers = static_cast(end - begin); + summary.has_run = has_run[row]; + summary.payload_bytes = + begin == end ? 0 : payload_offsets[end - 1] + payload_sizes[end - 1] - payload_offsets[begin]; + summary.invalid = valid_counts[row] != cardinality; +} + +__global__ void encode_batch_headers_kernel(std::uint32_t const* ids, + std::int64_t const* indptr, + std::int64_t const* valid_counts, + std::int64_t const* container_starts, + std::int64_t const* row_container_offsets, + container_kind const* kinds, + std::uint64_t const* payload_offsets, + device_build_summary const* summaries, + std::uint64_t const* output_offsets, + cuda::std::byte* storage) +{ + auto const row = static_cast(blockIdx.x); + auto const summary = summaries[row]; + if (summary.num_containers == 0) { return; } + auto const first_container = row_container_offsets[row]; + auto* output = storage + output_offsets[row]; + auto const has_run = summary.has_run != 0; + auto const run_bitmap_bytes = has_run ? (summary.num_containers + 7) / 8 : 0; + auto const descriptor_offset = + has_run ? sizeof(std::uint32_t) + run_bitmap_bytes : 2 * sizeof(std::uint32_t); + auto const offsets_offset = + descriptor_offset + summary.num_containers * 2 * sizeof(std::uint16_t); + auto const header_size = portable_header_size(summary.num_containers, has_run); + bool const store_offsets = !has_run || summary.num_containers >= kOffsetThreshold; + if (threadIdx.x == 0) { + if (has_run) { + write_u32(output, 0, kCookieRun | ((summary.num_containers - 1) << 16)); + } else { + write_u32(output, 0, kCookieNoRun); + write_u32(output, sizeof(std::uint32_t), summary.num_containers); + } + } + auto* output_bytes = reinterpret_cast(output); + for (std::uint32_t byte = threadIdx.x; byte < run_bitmap_bytes; byte += blockDim.x) { + std::uint8_t value{}; + for (std::uint32_t bit = 0; bit < 8; ++bit) { + auto const local = byte * 8 + bit; + if (local < summary.num_containers && kinds[first_container + local] == container_kind::run) { + value |= static_cast(1u << bit); + } + } + output_bytes[sizeof(std::uint32_t) + byte] = value; + } + auto const row_end = indptr[row] + valid_counts[row]; + for (std::uint32_t local = threadIdx.x; local < summary.num_containers; local += blockDim.x) { + auto const container = first_container + local; + auto const begin = container_starts[container]; + auto const end = local + 1 < summary.num_containers ? container_starts[container + 1] : row_end; + auto const descriptor = descriptor_offset + local * 2 * sizeof(std::uint16_t); + write_u16(output, descriptor, static_cast(ids[begin] >> 16)); + write_u16( + output, descriptor + sizeof(std::uint16_t), static_cast(end - begin - 1)); + if (store_offsets) { + auto const local_payload = payload_offsets[container] - payload_offsets[first_container]; + write_u32(output, + offsets_offset + local * sizeof(std::uint32_t), + static_cast(header_size + local_payload)); + } + } +} + +__global__ void encode_batch_payloads_kernel(std::uint32_t const* ids, + std::int64_t const* indptr, + std::int64_t const* valid_counts, + std::int64_t const* container_starts, + std::int64_t const* row_container_offsets, + std::uint32_t const* container_rows, + std::int64_t num_containers, + container_kind const* kinds, + std::uint64_t const* payload_sizes, + std::uint64_t const* payload_offsets, + device_build_summary const* summaries, + std::uint64_t const* output_offsets, + cuda::std::byte* storage) +{ + auto const container = static_cast(blockIdx.x); + if (container >= num_containers) { return; } + auto const row = static_cast(container_rows[container]); + auto const begin = container_starts[container]; + auto const first_container = row_container_offsets[row]; + auto const local = container - first_container; + auto const summary = summaries[row]; + auto const row_end = indptr[row] + valid_counts[row]; + auto const end = local + 1 < summary.num_containers ? container_starts[container + 1] : row_end; + auto* output = storage + output_offsets[row]; + auto const payload = portable_header_size(summary.num_containers, summary.has_run != 0) + + payload_offsets[container] - payload_offsets[first_container]; + + if (kinds[container] == container_kind::array) { + for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { + write_u16(output, + payload + static_cast(i - begin) * sizeof(std::uint16_t), + static_cast(ids[i] & 0xffffu)); + } + return; + } + using run_scan = cub::BlockScan; + union payload_scratch { + std::uint64_t bitmap_words[kBitmapBytes / sizeof(std::uint64_t)]; + typename run_scan::TempStorage run_scan_storage; + }; + __shared__ payload_scratch scratch; + __shared__ std::uint32_t run_base; + __shared__ std::uint32_t tile_runs; + if (kinds[container] == container_kind::bitmap) { + constexpr std::uint32_t words = kBitmapBytes / sizeof(std::uint64_t); + for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { + scratch.bitmap_words[word] = 0; + } + __syncthreads(); + for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { + auto const lower = ids[i] & 0xffffu; + atomicOr(reinterpret_cast(&scratch.bitmap_words[lower / 64]), + static_cast(std::uint64_t{1} << (lower % 64))); + } + __syncthreads(); + for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { + write_u64(output, payload + word * sizeof(std::uint64_t), scratch.bitmap_words[word]); + } + return; + } + auto const num_runs = static_cast( + (payload_sizes[container] - sizeof(std::uint16_t)) / (2 * sizeof(std::uint16_t))); + if (threadIdx.x == 0) { + write_u16(output, payload, num_runs); + run_base = 0; + } + __syncthreads(); + for (auto tile = begin; tile < end; tile += blockDim.x) { + auto const i = tile + threadIdx.x; + std::uint32_t const is_run_start = + i < end && (i == begin || + static_cast(ids[i]) != static_cast(ids[i - 1]) + 1) + ? 1u + : 0u; + std::uint32_t run_rank{}; + std::uint32_t block_runs{}; + run_scan(scratch.run_scan_storage).ExclusiveSum(is_run_start, run_rank, block_runs); + if (threadIdx.x == 0) { tile_runs = block_runs; } + __syncthreads(); + if (is_run_start != 0) { + auto j = i + 1; + while (j < end && + static_cast(ids[j]) == static_cast(ids[j - 1]) + 1) { + ++j; + } + auto const start = static_cast(ids[i] & 0xffffu); + auto const last = static_cast(ids[j - 1] & 0xffffu); + auto const run_offset = + payload + sizeof(std::uint16_t) + + static_cast(run_base + run_rank) * 2 * sizeof(std::uint16_t); + write_u16(output, run_offset, start); + write_u16( + output, run_offset + sizeof(std::uint16_t), static_cast(last - start)); + } + __syncthreads(); + if (threadIdx.x == 0) { run_base += tile_runs; } + __syncthreads(); + } +} + +__global__ void initialize_batch_refs_kernel(cuda::std::byte const* storage, + std::uint64_t const* row_offsets, + device_build_summary const* summaries, + ref_type* references, + std::size_t rows) +{ + auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (row < rows && summaries[row].num_containers != 0) { + ::new (static_cast(references + row)) ref_type{storage + row_offsets[row]}; + } +} + +batched_device_build_result build_general_rows( + raft::resources const& res, + std::size_t dataset_rows, + raft::device_vector_view ids, + raft::device_vector_view indptr, + std::vector const& host_indptr, + bool pre_sorted) +{ + common::nvtx::range build_scope( + "roaring_allowlist::build_general_batch"); + auto const stream = raft::resource::get_cuda_stream(res); + auto const rows = host_indptr.size() - 1; + auto const total_ids = static_cast(ids.extent(0)); + auto const item_count = static_cast(total_ids); + auto const row_count = static_cast(rows); + auto const chunks = + (static_cast(dataset_rows) + (std::uint64_t{1} << 16) - 1) >> 16; + std::size_t max_containers{}; + for (std::size_t row = 0; row < rows; ++row) { + auto const width = static_cast(host_indptr[row + 1] - host_indptr[row]); + max_containers += std::min(width, static_cast(chunks)); + } + RAFT_EXPECTS(rows <= std::numeric_limits::max(), + "Batched Roaring construction has too many rows for one launch."); + RAFT_EXPECTS(max_containers <= std::numeric_limits::max(), + "Batched Roaring construction has too many containers for one launch."); + constexpr int sort_end_bit = std::numeric_limits::digits; + auto counting = thrust::make_counting_iterator(0); + + std::size_t sort_workspace_bytes{}; + std::size_t select_workspace_bytes{}; + std::size_t scan_workspace_bytes{}; + if (!pre_sorted) { + RAFT_CUDA_TRY(cub::DeviceSegmentedRadixSort::SortKeys(nullptr, + sort_workspace_bytes, + ids.data_handle(), + static_cast(nullptr), + item_count, + row_count, + indptr.data_handle(), + indptr.data_handle() + 1, + 0, + sort_end_bit, + stream)); + } + RAFT_CUDA_TRY(cub::DeviceSelect::Flagged(nullptr, + select_workspace_bytes, + counting, + static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), + item_count, + stream)); + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, + scan_workspace_bytes, + static_cast(nullptr), + static_cast(nullptr), + static_cast(max_containers), + stream)); + auto const workspace_bytes = + std::max(sort_workspace_bytes, std::max(select_workspace_bytes, scan_workspace_bytes)); + batch_general_scratch_layout const layout{ + total_ids, rows, max_containers, !pre_sorted, workspace_bytes}; + rmm::device_uvector scratch(layout.bytes, stream); + auto* sorted_ids = + pre_sorted ? nullptr + : reinterpret_cast(scratch.data() + layout.sorted_ids_offset); + auto* valid_counts = reinterpret_cast(scratch.data() + layout.valid_counts_offset); + auto* flags = reinterpret_cast(scratch.data() + layout.flags_offset); + auto* selected_count = + reinterpret_cast(scratch.data() + layout.selected_count_offset); + auto* container_starts = + reinterpret_cast(scratch.data() + layout.container_starts_offset); + auto* row_container_offsets = + reinterpret_cast(scratch.data() + layout.row_container_offsets_offset); + auto* container_rows = + reinterpret_cast(scratch.data() + layout.container_rows_offset); + auto* kinds = reinterpret_cast(scratch.data() + layout.kinds_offset); + auto* payload_sizes = + reinterpret_cast(scratch.data() + layout.payload_sizes_offset); + auto* payload_offsets = + reinterpret_cast(scratch.data() + layout.payload_offsets_offset); + auto* has_run = reinterpret_cast(scratch.data() + layout.has_run_offset); + auto* summaries = + reinterpret_cast(scratch.data() + layout.summaries_offset); + auto* output_offsets = + reinterpret_cast(scratch.data() + layout.output_offsets_offset); + auto* workspace = scratch.data() + layout.workspace_offset; + + if (!pre_sorted) { + RAFT_CUDA_TRY(cub::DeviceSegmentedRadixSort::SortKeys(workspace, + sort_workspace_bytes, + ids.data_handle(), + sorted_ids, + item_count, + row_count, + indptr.data_handle(), + indptr.data_handle() + 1, + 0, + sort_end_bit, + stream)); + } + auto const* normalized = pre_sorted ? ids.data_handle() : sorted_ids; + find_batch_valid_counts_kernel<<>>( + normalized, indptr.data_handle(), row_count, dataset_rows, valid_counts); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + mark_batch_container_starts_kernel<<(rows), + kBuilderBlockSize, + 0, + stream>>>( + normalized, indptr.data_handle(), valid_counts, row_count, flags); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + RAFT_CUDA_TRY(cub::DeviceSelect::Flagged(workspace, + select_workspace_bytes, + counting, + flags, + container_starts, + selected_count, + item_count, + stream)); + find_row_container_offsets_kernel<<>>( + container_starts, selected_count, indptr.data_handle(), row_count, row_container_offsets); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + fill_container_rows_kernel<<(rows), kBuilderBlockSize, 0, stream>>>( + row_container_offsets, row_count, container_rows); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + RAFT_CUDA_TRY(cudaMemsetAsync(payload_sizes, 0, max_containers * sizeof(std::uint64_t), stream)); + RAFT_CUDA_TRY(cudaMemsetAsync(has_run, 0, rows * sizeof(std::uint32_t), stream)); + analyze_batch_containers_kernel<<(max_containers), + kBuilderBlockSize, + 0, + stream>>>(normalized, + indptr.data_handle(), + valid_counts, + container_starts, + selected_count, + container_rows, + kinds, + payload_sizes, + has_run); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(workspace, + scan_workspace_bytes, + payload_sizes, + payload_offsets, + static_cast(max_containers), + stream)); + finish_batch_rows_kernel<<>>( + row_count, + indptr.data_handle(), + valid_counts, + row_container_offsets, + has_run, + payload_sizes, + payload_offsets, + summaries); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + std::vector host_summaries(rows); + RAFT_CUDA_TRY(cudaMemcpyAsync(host_summaries.data(), + summaries, + rows * sizeof(device_build_summary), + cudaMemcpyDeviceToHost, + stream)); + raft::resource::sync_stream(res); + std::vector serialized(rows); + std::vector cardinalities(rows); + std::size_t actual_containers{}; + for (std::size_t row = 0; row < rows; ++row) { + auto const& summary = host_summaries[row]; + auto const cardinality = static_cast(host_indptr[row + 1] - host_indptr[row]); + RAFT_EXPECTS(summary.invalid == 0, "Roaring allowlist ID must be smaller than dataset_rows."); + RAFT_EXPECTS(summary.cardinality == static_cast(cardinality) && + (cardinality == 0 || summary.num_containers > 0), + "Internal error: general batched row analysis failed."); + cardinalities[row] = cardinality; + if (cardinality == 0) { continue; } + auto const header = portable_header_size(summary.num_containers, summary.has_run != 0); + RAFT_EXPECTS(summary.payload_bytes <= std::numeric_limits::max() - header, + "Portable Roaring row exceeds the 32-bit offset range."); + serialized[row] = header + static_cast(summary.payload_bytes); + actual_containers += summary.num_containers; + } + auto packed = make_packed_rows_layout(serialized); + rmm::device_uvector storage(packed.bytes, stream); + static_assert(sizeof(std::size_t) == sizeof(std::uint64_t)); + RAFT_CUDA_TRY(cudaMemcpyAsync(output_offsets, + packed.row_offsets.data(), + rows * sizeof(std::uint64_t), + cudaMemcpyHostToDevice, + stream)); + encode_batch_headers_kernel<<(rows), kBuilderBlockSize, 0, stream>>>( + normalized, + indptr.data_handle(), + valid_counts, + container_starts, + row_container_offsets, + kinds, + payload_offsets, + summaries, + output_offsets, + storage.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + encode_batch_payloads_kernel<<(actual_containers), + kBuilderBlockSize, + 0, + stream>>>(normalized, + indptr.data_handle(), + valid_counts, + container_starts, + row_container_offsets, + container_rows, + static_cast(actual_containers), + kinds, + payload_sizes, + payload_offsets, + summaries, + output_offsets, + storage.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + initialize_batch_refs_kernel<<>>( + storage.data(), + output_offsets, + summaries, + reinterpret_cast(storage.data() + packed.references_offset), + rows); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + return {std::move(storage), + std::move(packed.row_offsets), + std::move(packed.reference_offsets), + std::move(serialized), + std::move(cardinalities), + true}; +} + +void validate_indptr(std::vector const& indptr, std::size_t rows, std::size_t nnz) +{ + RAFT_EXPECTS(rows > 0, "Roaring input must contain at least one allowlist row."); + RAFT_EXPECTS(indptr.size() == rows + 1, "Roaring indptr must contain num_rows + 1 entries."); + RAFT_EXPECTS(indptr.front() == 0, "Roaring indptr must start at zero."); + for (std::size_t row = 0; row < rows; ++row) { + RAFT_EXPECTS(indptr[row] <= indptr[row + 1] && indptr[row] >= 0, + "Roaring indptr must be nonnegative and nondecreasing."); + } + RAFT_EXPECTS(indptr.back() >= 0 && static_cast(indptr.back()) == nnz, + "The final Roaring indptr entry must equal nnz."); +} + +batched_device_build_result build_batched_from_device_ids( + raft::resources const& res, + std::size_t dataset_rows, + raft::device_vector_view ids, + raft::device_vector_view indptr, + std::vector const& host_indptr, + bool pre_sorted) +{ + auto const stream = raft::resource::get_cuda_stream(res); + auto const rows = host_indptr.size() - 1; + auto const size = static_cast(ids.extent(0)); + validate_indptr(host_indptr, rows, size); + if (size == 0) { + return {rmm::device_uvector(0, stream), + std::vector(rows), + std::vector(rows), + std::vector(rows), + std::vector(rows), + true}; + } + if (rows == 1) { + auto row_view = raft::make_device_vector_view( + ids.data_handle(), static_cast(size)); + auto built = build_from_device_ids(res, dataset_rows, row_view, pre_sorted); + return {std::move(built.storage), + {0}, + {built.cardinality == 0 ? 0 : reference_offset(built.serialized_bytes)}, + {built.serialized_bytes}, + {built.cardinality}, + built.reference_initialized}; + } + return build_general_rows(res, dataset_rows, ids, indptr, host_indptr, pre_sorted); +} + +/** + * Validate one externally supplied portable row before cuco sees it. + * + * cuco's raw-byte reference constructor assumes a valid stream and is not given + * the row's byte length. We therefore walk the complete row on the host first, + * checking every read, the required header variant, ordered keys and values, + * payload cardinalities, exact container offsets, the logical dataset bound, + * and that no trailing bytes remain. Besides rejecting malformed input, this + * walk records cardinality without retaining another decoded representation. + * + * A zero-length input is a cuVS convenience for an empty allowlist. The + * standard serialized empty form (no-run cookie followed by N = 0) is accepted + * as well. + */ +row_metadata validate_serialized_row(std::byte const* data, + std::size_t size, + std::size_t dataset_rows) +{ + // Empty here means the outer cuVS byte offsets selected no portable bytes for + // this row. + if (size == 0) { return {}; } + + // Decode the cookie first because it determines both how N is stored and + // whether a run bitmap follows. Every subsequent read advances `cursor` + // through exactly one format field. + auto const cookie = read_u32(data, size, 0); + bool const has_run = (cookie & 0xffffu) == kCookieRun; + std::size_t num_containers{}; + std::size_t cursor = 4; + if (has_run) { + num_containers = (cookie >> 16) + 1; + } else { + RAFT_EXPECTS(cookie == kCookieNoRun, "Malformed portable Roaring bitmap: unsupported cookie."); + num_containers = read_u32(data, size, cursor); + cursor += 4; + } + RAFT_EXPECTS(num_containers <= (std::size_t{1} << 16), + "Malformed portable Roaring bitmap: too many containers."); + if (num_containers == 0) { + RAFT_EXPECTS(!has_run && cursor == size, + "Malformed portable Roaring bitmap: invalid empty representation."); + return {}; + } + + std::size_t run_bitmap_offset{}; + if (has_run) { + auto const run_bitmap_bytes = (num_containers + 7) / 8; + RAFT_EXPECTS(cursor <= size && size - cursor >= run_bitmap_bytes, + "Malformed portable Roaring bitmap: truncated run bitmap."); + run_bitmap_offset = cursor; + cursor += run_bitmap_bytes; + } + + // The descriptive header is common to both cookie forms. Reconstruct + // cardinality by adding one to its encoded value and require container keys + // to be strictly increasing. + std::vector keys(num_containers); + std::vector cards(num_containers); + for (std::size_t i = 0; i < num_containers; ++i) { + keys[i] = read_u16(data, size, cursor); + cards[i] = static_cast(read_u16(data, size, cursor + 2)) + 1; + cursor += 4; + if (i > 0) { + RAFT_EXPECTS(keys[i - 1] < keys[i], + "Malformed portable Roaring bitmap: container keys are not ordered."); + } + } + + // These are offsets INSIDE this portable row. The no-run form always stores + // them; the run form stores them only at the specification's four-container + // threshold. + auto const store_offsets = !has_run || num_containers >= kOffsetThreshold; + std::vector container_offsets; + if (store_offsets) { + container_offsets.resize(num_containers); + for (std::size_t i = 0; i < num_containers; ++i) { + container_offsets[i] = read_u32(data, size, cursor); + cursor += 4; + } + } + + // Validate payloads in descriptor order. Requiring every stored offset to + // equal `cursor` also rejects gaps, overlaps, and offsets that point into a + // header or a different container. + row_metadata metadata; + metadata.empty = false; + for (std::size_t i = 0; i < num_containers; ++i) { + if (store_offsets) { + RAFT_EXPECTS(container_offsets[i] == cursor, + "Malformed portable Roaring bitmap: invalid container offset."); + } + auto const is_run = + has_run && ((std::to_integer(data[run_bitmap_offset + i / 8]) >> (i % 8)) & 1u); + std::uint32_t lower_max{}; + if (is_run) { + auto const num_runs = read_u16(data, size, cursor); + cursor += 2; + std::uint32_t run_cardinality{}; + std::uint32_t previous_end{}; + for (std::size_t run_index = 0; run_index < num_runs; ++run_index) { + auto const start = static_cast(read_u16(data, size, cursor)); + auto const length = static_cast(read_u16(data, size, cursor + 2)); + cursor += 4; + auto const end = start + length; + RAFT_EXPECTS(end <= std::numeric_limits::max(), + "Malformed portable Roaring bitmap: run exceeds uint16 range."); + if (run_index > 0) { + RAFT_EXPECTS(start > previous_end, + "Malformed portable Roaring bitmap: runs overlap or are " + "unordered."); + } + previous_end = end; + lower_max = end; + run_cardinality += length + 1; + } + RAFT_EXPECTS(num_runs > 0 && run_cardinality == cards[i], + "Malformed portable Roaring bitmap: invalid run cardinality."); + } else if (cards[i] <= kArrayCardinality) { + std::uint32_t previous{}; + for (std::size_t j = 0; j < cards[i]; ++j) { + auto const value = static_cast(read_u16(data, size, cursor)); + cursor += 2; + if (j > 0) { + RAFT_EXPECTS(previous < value, + "Malformed portable Roaring bitmap: " + "array values are not ordered."); + } + previous = value; + lower_max = value; + } + } else { + RAFT_EXPECTS(cursor <= size && size - cursor >= kBitmapBytes, + "Malformed portable Roaring bitmap: truncated bitmap container."); + std::size_t popcount{}; + bool found_max = false; + for (std::size_t j = 0; j < kBitmapBytes; ++j) { + popcount += std::popcount(std::to_integer(data[cursor + j])); + } + for (std::size_t j = kBitmapBytes; j-- > 0 && !found_max;) { + auto const byte = std::to_integer(data[cursor + j]); + if (byte != 0) { + lower_max = static_cast(j * 8 + (7 - std::countl_zero(byte))); + found_max = true; + } + } + RAFT_EXPECTS(found_max && popcount == cards[i], + "Malformed portable Roaring bitmap: invalid bitmap cardinality."); + cursor += kBitmapBytes; + } + + auto const max_id = (static_cast(keys[i]) << 16) | lower_max; + RAFT_EXPECTS(static_cast(max_id) < dataset_rows, + "Portable Roaring bitmap contains an ID outside dataset_rows."); + metadata.cardinality += cards[i]; + metadata.max_id = max_id; + } + RAFT_EXPECTS(cursor == size, "Malformed portable Roaring bitmap: trailing or unconsumed bytes."); + return metadata; +} + +/** + * Construct the lightweight cuco reference once, outside the search path. + * + * The raw-byte constructor parses the portable header and stores small + * container-location metadata by value while retaining pointers into `data`. It + * neither allocates nor copies the serialized payload. Empty allowlists skip + * this kernel because the pinned parser expects at least one container. + */ +__global__ void initialize_imported_refs_kernel(cuda::std::byte const* storage, + std::uint64_t const* row_offsets, + std::uint64_t const* serialized_bytes, + ref_type* references, + std::size_t rows) +{ + auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (row < rows && serialized_bytes[row] != 0) { + ::new (static_cast(references + row)) ref_type{storage + row_offsets[row]}; + } +} + +__global__ void initialize_view_tables_kernel(cuda::std::byte const* storage, + std::uint64_t const* reference_offsets, + std::uint64_t const* serialized_bytes, + ref_type const** references, + std::uint8_t* empty_rows, + std::size_t rows) +{ + auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (row >= rows) { return; } + auto const empty = serialized_bytes[row] == 0; + references[row] = + empty ? nullptr : reinterpret_cast(storage + reference_offsets[row]); + empty_rows[row] = empty ? 1 : 0; +} + +__global__ void contains_kernel(ref_type const* const* references, + std::uint8_t const* empty_rows, + std::uint64_t dataset_rows, + std::uint32_t const* row_ids, + std::uint8_t* output, + std::size_t columns, + std::size_t size) +{ + auto const i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= size) { return; } + auto const query = i / columns; + auto const row = row_ids[i]; + output[i] = empty_rows[query] == 0 && static_cast(row) < dataset_rows && + references[query]->contains(row); +} + +} // namespace + +struct roaring_allowlist::impl { + rmm::device_uvector storage; + std::vector row_offsets_; + std::vector reference_offsets_; + std::vector serialized_bytes_; + std::vector cardinalities_; + rmm::device_uvector references; + rmm::device_uvector empty_rows; + std::size_t dataset_rows_{}; + std::size_t total_cardinality_{}; + bool references_initialized_{}; + + impl(raft::resources const& res, std::size_t dataset_rows, batched_device_build_result&& built) + : storage(std::move(built.storage)), + row_offsets_(std::move(built.row_offsets)), + reference_offsets_(std::move(built.reference_offsets)), + serialized_bytes_(std::move(built.serialized_bytes)), + cardinalities_(std::move(built.cardinalities)), + references(cardinalities_.size(), raft::resource::get_cuda_stream(res)), + empty_rows(cardinalities_.size(), raft::resource::get_cuda_stream(res)), + dataset_rows_(dataset_rows), + total_cardinality_( + std::accumulate(cardinalities_.begin(), cardinalities_.end(), std::size_t{})), + references_initialized_(built.references_initialized) + { + static_assert(std::is_trivially_destructible_v); + auto const rows = cardinalities_.size(); + RAFT_EXPECTS(rows > 0 && row_offsets_.size() == rows && reference_offsets_.size() == rows && + serialized_bytes_.size() == rows, + "Internal error: inconsistent batched Roaring metadata."); + if (rows == 0) { return; } + + auto const stream = raft::resource::get_cuda_stream(res); + static_assert(sizeof(std::size_t) == sizeof(std::uint64_t)); + rmm::device_uvector device_row_offsets(rows, stream); + rmm::device_uvector device_reference_offsets(rows, stream); + rmm::device_uvector device_serialized_bytes(rows, stream); + RAFT_CUDA_TRY(cudaMemcpyAsync(device_row_offsets.data(), + row_offsets_.data(), + rows * sizeof(std::uint64_t), + cudaMemcpyHostToDevice, + stream)); + RAFT_CUDA_TRY(cudaMemcpyAsync(device_reference_offsets.data(), + reference_offsets_.data(), + rows * sizeof(std::uint64_t), + cudaMemcpyHostToDevice, + stream)); + RAFT_CUDA_TRY(cudaMemcpyAsync(device_serialized_bytes.data(), + serialized_bytes_.data(), + rows * sizeof(std::uint64_t), + cudaMemcpyHostToDevice, + stream)); + if (!references_initialized_ && storage.size() != 0) { + initialize_imported_refs_kernel<<>>( + storage.data(), + device_row_offsets.data(), + device_serialized_bytes.data(), + reinterpret_cast(storage.data() + reference_offsets_.front()), + rows); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + initialize_view_tables_kernel<<>>( + storage.data(), + device_reference_offsets.data(), + device_serialized_bytes.data(), + references.data(), + empty_rows.data(), + rows); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + + [[nodiscard]] ref_type const* reference(std::size_t allowlist_id) const noexcept + { + return cardinalities_[allowlist_id] == 0 + ? nullptr + : reinterpret_cast(storage.data() + reference_offsets_[allowlist_id]); + } +}; + +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, + raft::host_vector_view indptr, + bool pre_sorted) +{ + validate_dataset_rows(dataset_rows); + RAFT_EXPECTS(indptr.extent(0) >= 2, "Roaring indptr must contain at least two entries."); + auto const rows = static_cast(indptr.extent(0) - 1); + auto const size = static_cast(ids.extent(0)); + std::vector host_indptr(indptr.data_handle(), + indptr.data_handle() + indptr.extent(0)); + validate_indptr(host_indptr, rows, size); + + auto const stream = raft::resource::get_cuda_stream(res); + rmm::device_uvector device_ids(size, stream); + rmm::device_uvector device_indptr(rows + 1, stream); + if (size != 0) { + RAFT_CUDA_TRY(cudaMemcpyAsync(device_ids.data(), + ids.data_handle(), + size * sizeof(key_type), + cudaMemcpyHostToDevice, + stream)); + RAFT_CUDA_TRY(cudaMemcpyAsync(device_indptr.data(), + host_indptr.data(), + (rows + 1) * sizeof(indptr_type), + cudaMemcpyHostToDevice, + stream)); + } + auto device_ids_view = raft::make_device_vector_view( + device_ids.data(), static_cast(size)); + auto device_indptr_view = raft::make_device_vector_view( + device_indptr.data(), static_cast(rows + 1)); + auto built = build_batched_from_device_ids( + res, dataset_rows, device_ids_view, device_indptr_view, host_indptr, 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, + raft::device_vector_view indptr, + bool pre_sorted) +{ + validate_dataset_rows(dataset_rows); + RAFT_EXPECTS(indptr.extent(0) >= 2, "Roaring indptr must contain at least two entries."); + auto const rows = static_cast(indptr.extent(0) - 1); + auto const size = static_cast(ids.extent(0)); + auto const stream = raft::resource::get_cuda_stream(res); + std::vector host_indptr(rows + 1); + RAFT_CUDA_TRY(cudaMemcpyAsync(host_indptr.data(), + indptr.data_handle(), + (rows + 1) * sizeof(indptr_type), + cudaMemcpyDeviceToHost, + stream)); + raft::resource::sync_stream(res); + validate_indptr(host_indptr, rows, size); + auto built = + build_batched_from_device_ids(res, dataset_rows, ids, indptr, host_indptr, pre_sorted); + return roaring_allowlist{std::make_unique(res, dataset_rows, std::move(built))}; +} + +roaring_allowlist roaring_allowlist::from_serialized( + raft::resources const& res, + std::size_t dataset_rows, + raft::host_vector_view bytes, + raft::host_vector_view byte_offsets) +{ + validate_dataset_rows(dataset_rows); + RAFT_EXPECTS(byte_offsets.extent(0) >= 2, + "Roaring byte_offsets must contain at least two entries."); + auto const rows = static_cast(byte_offsets.extent(0) - 1); + auto const size = static_cast(bytes.extent(0)); + RAFT_EXPECTS(byte_offsets(0) == 0, "Roaring byte_offsets must start at zero."); + RAFT_EXPECTS(byte_offsets(static_cast(rows)) == size, + "The final Roaring byte offset must equal bytes.extent(0)."); + + std::vector serialized_bytes(rows); + std::vector cardinalities(rows); + for (std::size_t row = 0; row < rows; ++row) { + auto const begin = byte_offsets(static_cast(row)); + auto const end = byte_offsets(static_cast(row + 1)); + RAFT_EXPECTS(begin <= end && end <= size, + "Roaring byte_offsets must be nondecreasing and in bounds."); + auto const row_size = static_cast(end - begin); + auto const metadata = + validate_serialized_row(bytes.data_handle() + begin, row_size, dataset_rows); + serialized_bytes[row] = metadata.empty ? 0 : row_size; + cardinalities[row] = metadata.cardinality; + } + + auto const packed = make_packed_rows_layout(serialized_bytes); + auto const stream = raft::resource::get_cuda_stream(res); + rmm::device_uvector storage(packed.bytes, stream); + for (std::size_t row = 0; row < rows; ++row) { + if (serialized_bytes[row] == 0) { continue; } + auto const begin = byte_offsets(static_cast(row)); + RAFT_CUDA_TRY(cudaMemcpyAsync(storage.data() + packed.row_offsets[row], + bytes.data_handle() + begin, + serialized_bytes[row], + cudaMemcpyHostToDevice, + stream)); + } + + batched_device_build_result built{std::move(storage), + packed.row_offsets, + packed.reference_offsets, + std::move(serialized_bytes), + std::move(cardinalities), + false}; + 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::num_allowlists() const noexcept +{ + return impl_->cardinalities_.size(); +} + +std::size_t roaring_allowlist::dataset_rows() const noexcept { return impl_->dataset_rows_; } + +std::size_t roaring_allowlist::cardinality(std::size_t allowlist_id) const +{ + RAFT_EXPECTS(allowlist_id < num_allowlists(), "Roaring allowlist_id is out of range."); + return impl_->cardinalities_[allowlist_id]; +} + +bool roaring_allowlist::empty(std::size_t allowlist_id) const +{ + return cardinality(allowlist_id) == 0; +} + +std::size_t roaring_allowlist::total_cardinality() const noexcept +{ + return impl_->total_cardinality_; +} + +std::size_t roaring_allowlist::size_bytes() const noexcept +{ + return impl_->storage.size() * sizeof(cuda::std::byte) + + impl_->references.size() * sizeof(ref_type const*) + + impl_->empty_rows.size() * sizeof(std::uint8_t); +} + +roaring_allowlist_view roaring_allowlist::view(std::size_t allowlist_id) const +{ + RAFT_EXPECTS(allowlist_id < num_allowlists(), "Roaring allowlist_id is out of range."); + return roaring_allowlist_view{ + impl_->reference(allowlist_id), dataset_rows(), cardinality(allowlist_id)}; +} + +void roaring_allowlist::contains( + raft::resources const& res, + raft::device_matrix_view row_ids, + raft::device_matrix_view output) const +{ + contains_async(res, row_ids, output); + raft::resource::sync_stream(res); +} + +void roaring_allowlist::contains_async( + raft::resources const& res, + raft::device_matrix_view row_ids, + raft::device_matrix_view output) const +{ + RAFT_EXPECTS(row_ids.extent(0) == static_cast(num_allowlists()), + "Roaring membership matrix must have one row per allowlist."); + RAFT_EXPECTS(output.extent(0) == row_ids.extent(0) && output.extent(1) == row_ids.extent(1), + "Roaring membership output shape must match the input shape."); + auto const rows = static_cast(row_ids.extent(0)); + auto const columns = static_cast(row_ids.extent(1)); + if (columns == 0) { return; } + RAFT_EXPECTS(rows <= std::numeric_limits::max() / columns, + "Roaring membership matrix is too large."); + auto const size = rows * columns; + constexpr int block_size = 256; + contains_kernel<<>>( + impl_->references.data(), + impl_->empty_rows.data(), + static_cast(dataset_rows()), + row_ids.data_handle(), + output.data_handle(), + columns, + size); + RAFT_CUDA_TRY(cudaPeekAtLastError()); +} + +} // namespace cuvs::core diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 80e2f2a07e..b040cece7c 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -20,6 +20,7 @@ #include #include +#include #include #include @@ -442,6 +443,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 3fd1963268..03278b8e81 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -70,6 +70,8 @@ cuvs::neighbors::cagra::index merge( "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 7cad56767a..8feaf7fcd3 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/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/roaring_filter.cu b/cpp/src/neighbors/roaring_filter.cu new file mode 100644 index 0000000000..148186443f --- /dev/null +++ b/cpp/src/neighbors/roaring_filter.cu @@ -0,0 +1,190 @@ +/* + * 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 +#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) +{ + 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_CUDA_TRY(cudaMemcpyAsync(refs.data(), + host_refs.data(), + host_refs.size() * sizeof(ref_type const*), + cudaMemcpyHostToDevice, + stream)); + RAFT_CUDA_TRY(cudaMemcpyAsync(empty_rows.data(), + host_empty.data(), + host_empty.size() * sizeof(std::uint8_t), + cudaMemcpyHostToDevice, + stream)); + auto const host_payload = data_type{refs.data(), + empty_rows.data(), + static_cast(allowlists.size()), + static_cast(dataset_rows_)}; + RAFT_CUDA_TRY(cudaMemcpyAsync( + payload.data(), &host_payload, sizeof(host_payload), cudaMemcpyHostToDevice, 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_CUDA_TRY(cudaMemcpyAsync( + impl_->refs.data() + query_id, &ref, sizeof(ref), cudaMemcpyHostToDevice, stream)); + RAFT_CUDA_TRY(cudaMemcpyAsync( + impl_->empty_rows.data() + query_id, &empty_, sizeof(empty_), cudaMemcpyHostToDevice, 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 437396b736..b7a8c2c71e 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -185,6 +185,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..e9f7534bea 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 let the same fixture exercise rectangular batched construction. + 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,115 @@ TEST_P(CagraUdfFilterTest, TenantContextHonorsQuerySpecificMetadata) EXPECT_EQ(host_row_tenants[source_id], query_tenant); } } + + // Reuse this existing query-specific UDF test as the exact reference for every single-partition + // algorithm, including max_queries=2 chunking above. + std::vector allowed_ids; + std::vector indptr{0}; + for (std::int64_t q = 0; q < n_queries; ++q) { + auto query_tenant = host_query_tenants[static_cast(q)]; + 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)); + } + } + indptr.push_back(static_cast(allowed_ids.size())); + } + + auto const tenant_width = indptr[1] - indptr[0]; + ASSERT_GT(tenant_width, 0); + for (std::int64_t query = 0; query < n_queries; ++query) { + ASSERT_EQ(indptr[static_cast(query + 1)] - indptr[static_cast(query)], + tenant_width); + } + auto tenant_allowlists = cuvs::core::roaring_allowlist::from_ids( + res, + n_rows, + raft::make_host_vector_view(allowed_ids.data(), + allowed_ids.size()), + raft::make_host_vector_view(indptr.data(), indptr.size())); + std::vector tenant_views; + tenant_views.reserve(n_queries); + for (std::int64_t query = 0; query < n_queries; ++query) { + tenant_views.push_back(tenant_allowlists.view(query)); + } + 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); + + // Keep edge-selectivity coverage in this existing fixture: one filter mixes empty, full, + // fewer-than-k, sparse, and dense query allowlists. + std::vector> mixed_rows(static_cast(n_queries)); + for (std::uint32_t row = 0; row < n_rows; ++row) { + mixed_rows[1].push_back(row); + } + mixed_rows[2] = {7}; + mixed_rows[3] = {10, 11, 12}; + for (std::uint32_t row = 0; row < n_rows; row += 8) { + mixed_rows[4].push_back(row); + } + for (std::uint32_t row = n_rows - 32; row < n_rows; ++row) { + mixed_rows[5].push_back(row); + } + + std::vector mixed_ids; + std::vector mixed_indptr{0}; + for (auto const& row : mixed_rows) { + mixed_ids.insert(mixed_ids.end(), row.begin(), row.end()); + mixed_indptr.push_back(static_cast(mixed_ids.size())); + } + auto mixed_allowlists = cuvs::core::roaring_allowlist::from_ids( + res, + n_rows, + raft::make_host_vector_view(mixed_ids.data(), + mixed_ids.size()), + raft::make_host_vector_view(mixed_indptr.data(), + mixed_indptr.size()), + true); + std::vector mixed_views; + mixed_views.reserve(mixed_rows.size()); + for (std::size_t query = 0; query < mixed_rows.size(); ++query) { + mixed_views.push_back(mixed_allowlists.view(query)); + } + cuvs::neighbors::filtering::roaring_filter mixed_filter(res, mixed_views); + auto mixed_result = search(mixed_filter); + for (std::int64_t query = 0; query < n_queries; ++query) { + auto const& allowed = mixed_rows[static_cast(query)]; + for (std::int64_t rank = 0; rank < k; ++rank) { + auto row = mixed_result.neighbors[static_cast(query * k + rank)]; + auto valid_row = row < static_cast(n_rows); + if (query == 0) { EXPECT_FALSE(valid_row); } + if (query == 1) { EXPECT_TRUE(valid_row); } + if (valid_row) { EXPECT_NE(std::find(allowed.begin(), allowed.end(), row), allowed.end()); } + } + } + + if (GetParam() == cagra::search_algo::SINGLE_CTA) { + std::array const one_empty{0, 0}; + auto wrong_queries = cuvs::core::roaring_allowlist::from_ids( + res, + n_rows, + raft::make_host_vector_view(nullptr, 0), + raft::make_host_vector_view(one_empty.data(), + one_empty.size())); + std::array wrong_query_views{wrong_queries.view(0)}; + cuvs::neighbors::filtering::roaring_filter wrong_query_filter(res, wrong_query_views); + EXPECT_THROW(search(wrong_query_filter), raft::logic_error); + + std::vector const all_empty(static_cast(n_queries) + 1, 0); + auto wrong_columns = cuvs::core::roaring_allowlist::from_ids( + res, + n_rows + 1, + raft::make_host_vector_view(nullptr, 0), + raft::make_host_vector_view(all_empty.data(), + all_empty.size())); + std::vector wrong_column_views; + for (std::int64_t query = 0; query < n_queries; ++query) { + wrong_column_views.push_back(wrong_columns.view(query)); + } + 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..b744af3935 --- /dev/null +++ b/cpp/tests/neighbors/roaring_allowlist.cu @@ -0,0 +1,944 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +#include "../../bench/ann/src/cuvs/allowlist_benchmark_utils.hpp" +#include "../../bench/ann/src/cuvs/cagra_filter_benchmark_utils.hpp" + +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::core { +namespace { + +// These hand-built fixtures exercise import of the portable bytes themselves. +// Keep their field order aligned with +// https://github.com/RoaringBitmap/RoaringFormatSpec. +constexpr std::uint32_t kCookieNoRun = 12346; +constexpr std::uint32_t kCookieRun = 12347; +constexpr std::uint32_t kSingleArrayPayloadByteOffset = 16; + +void append_u16(std::vector& out, std::uint16_t value) +{ + out.push_back(static_cast(value & 0xffu)); + out.push_back(static_cast((value >> 8) & 0xffu)); +} + +void append_u32(std::vector& out, std::uint32_t value) +{ + for (int i = 0; i < 4; ++i) { + out.push_back(static_cast((value >> (8 * i)) & 0xffu)); + } +} + +// One no-run array container: cookie, N, {key, cardinality - 1}, payload +// offset, values. +std::vector array_row(std::vector const& values) +{ + std::vector out; + append_u32(out, kCookieNoRun); + append_u32(out, 1); + append_u16(out, 0); + append_u16(out, static_cast(values.size() - 1)); + append_u32(out, kSingleArrayPayloadByteOffset); + for (auto value : values) { + append_u16(out, value); + } + return out; +} + +// One run container: run cookie, one-byte run bitmap, descriptor, then the run +// payload. A run-cookie row with fewer than four containers has no +// container-offset table. +std::vector run_row(std::uint16_t start, std::uint16_t length_minus_one) +{ + std::vector out; + append_u32(out, kCookieRun); + out.push_back(std::byte{1}); + append_u16(out, 0); + append_u16(out, length_minus_one); + append_u16(out, 1); + append_u16(out, start); + append_u16(out, length_minus_one); + return out; +} + +roaring_allowlist from_ragged_ids(raft::resources const& res, + std::size_t dataset_rows, + std::vector const& ids, + std::vector const& indptr, + bool pre_sorted = false) +{ + return roaring_allowlist::from_ids( + res, + dataset_rows, + raft::make_host_vector_view(ids.data(), ids.size()), + raft::make_host_vector_view(indptr.data(), indptr.size()), + pre_sorted); +} + +roaring_allowlist from_device_ragged_ids(raft::resources const& res, + std::size_t dataset_rows, + std::vector const& ids, + std::vector const& indptr, + bool pre_sorted = false) +{ + auto device_ids = raft::make_device_vector(res, ids.size()); + auto device_indptr = raft::make_device_vector(res, indptr.size()); + auto const stream = raft::resource::get_cuda_stream(res); + raft::update_device(device_ids.data_handle(), ids.data(), ids.size(), stream); + raft::update_device(device_indptr.data_handle(), indptr.data(), indptr.size(), stream); + return roaring_allowlist::from_ids( + res, + dataset_rows, + raft::make_device_vector_view(device_ids.data_handle(), + ids.size()), + raft::make_device_vector_view(device_indptr.data_handle(), + indptr.size()), + pre_sorted); +} + +roaring_allowlist from_ids(raft::resources const& res, + std::size_t dataset_rows, + std::vector const& ids, + bool pre_sorted = false) +{ + return from_ragged_ids( + res, dataset_rows, ids, {0, static_cast(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) +{ + return from_device_ragged_ids( + res, dataset_rows, ids, {0, static_cast(ids.size())}, pre_sorted); +} + +roaring_allowlist import_row(raft::resources const& res, + std::size_t dataset_rows, + std::vector const& bytes) +{ + std::array const offsets{0, bytes.size()}; + return roaring_allowlist::from_serialized( + res, + dataset_rows, + raft::make_host_vector_view( + bytes.data(), static_cast(bytes.size())), + raft::make_host_vector_view(offsets.data(), offsets.size())); +} + +std::vector copy_serialized_bytes(raft::resources const& res, + roaring_allowlist const& allowlist) +{ + using ref_type = cuco::experimental::roaring_bitmap_ref; + static_assert(std::is_trivially_copyable_v); + auto const view = allowlist.view(0); + EXPECT_FALSE(view.empty()); + EXPECT_NE(view.device_reference(), nullptr); + + std::array raw_reference{}; + auto const stream = raft::resource::get_cuda_stream(res); + RAFT_CUDA_TRY(cudaMemcpyAsync(raw_reference.data(), + view.device_reference(), + raw_reference.size(), + cudaMemcpyDeviceToHost, + stream)); + raft::resource::sync_stream(res); + auto const reference = std::bit_cast(raw_reference); + std::vector bytes(reference.size_bytes()); + RAFT_CUDA_TRY( + cudaMemcpyAsync(bytes.data(), reference.data(), bytes.size(), cudaMemcpyDeviceToHost, stream)); + raft::resource::sync_stream(res); + return bytes; +} + +void expect_membership(raft::resources const& res, + roaring_allowlist const& allowlist, + std::vector const& row_ids, + std::vector const& expected, + bool asynchronous = false) +{ + 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); + if (asynchronous) { + allowlist.contains_async( + res, + raft::make_device_matrix_view( + rows.data_handle(), 1, static_cast(row_ids.size())), + raft::make_device_matrix_view( + output.data_handle(), 1, static_cast(expected.size()))); + } else { + allowlist.contains( + res, + raft::make_device_matrix_view( + rows.data_handle(), 1, static_cast(row_ids.size())), + raft::make_device_matrix_view( + output.data_handle(), 1, static_cast(expected.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); +} + +void expect_batch_membership(raft::resources const& res, + roaring_allowlist const& allowlists, + std::size_t columns, + std::vector const& row_ids, + std::vector const& expected) +{ + ASSERT_EQ(row_ids.size(), expected.size()); + ASSERT_EQ(row_ids.size(), allowlists.num_allowlists() * columns); + auto device_rows = raft::make_device_vector(res, row_ids.size()); + auto output = raft::make_device_vector(res, expected.size()); + auto const stream = raft::resource::get_cuda_stream(res); + raft::update_device(device_rows.data_handle(), row_ids.data(), row_ids.size(), stream); + allowlists.contains( + res, + raft::make_device_matrix_view( + device_rows.data_handle(), + static_cast(allowlists.num_allowlists()), + static_cast(columns)), + raft::make_device_matrix_view( + output.data_handle(), + static_cast(allowlists.num_allowlists()), + static_cast(columns))); + 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, BuildsArrayRunBitmapAndMultiContainerRowsFromIds) +{ + raft::device_resources res; + + auto array = from_ids(res, 200000, {65537, 7, 3, 131074, 5}); + EXPECT_EQ(array.dataset_rows(), 200000); + EXPECT_EQ(array.cardinality(0), 5); + EXPECT_FALSE(array.empty(0)); + EXPECT_GT(array.size_bytes(), 0); + expect_membership(res, array, {3, 4, 7, 65537, 131073, 131074, 200000}, {1, 0, 1, 1, 0, 1, 0}); + + std::vector consecutive; + for (std::uint32_t id = 100; id < 300; ++id) { + consecutive.push_back(id); + } + auto run = from_ids(res, 1000, consecutive); + EXPECT_EQ(run.cardinality(0), 200); + expect_membership(res, run, {99, 100, 199, 299, 300}, {0, 1, 1, 1, 0}); + + std::vector sparse; + for (std::uint32_t id = 0; id < 10000; id += 2) { + sparse.push_back(id); + } + auto bitmap = from_ids(res, 10000, sparse); + EXPECT_EQ(bitmap.cardinality(0), 5000); + expect_membership(res, bitmap, {0, 1, 8192, 9998, 9999}, {1, 0, 1, 1, 0}, true); +} + +TEST(RoaringAllowlist, BuildsRaggedRowsInOneGeneralBatch) +{ + raft::device_resources res; + constexpr std::size_t rows = 4; + constexpr std::size_t dataset_rows = (std::size_t{64} << 16) + 16; + std::vector ids; + std::vector indptr{0}; + + for (std::uint32_t i = 0; i < 64; ++i) { + ids.push_back(i); + } + indptr.push_back(static_cast(ids.size())); + indptr.push_back(static_cast(ids.size())); // empty row + for (std::uint32_t i = 0; i < 128; ++i) { + ids.push_back(1000 + 2 * i); + } + indptr.push_back(static_cast(ids.size())); + for (std::uint32_t i = 0; i < 64; ++i) { + ids.push_back((i << 16) + 7); + } + indptr.push_back(static_cast(ids.size())); + for (std::size_t row = 0; row < rows; ++row) { + std::reverse(ids.begin() + indptr[row], ids.begin() + indptr[row + 1]); + } + + auto allowlists = from_ragged_ids(res, dataset_rows, ids, indptr); + auto device_allowlists = from_device_ragged_ids(res, dataset_rows, ids, indptr); + EXPECT_EQ(allowlists.num_allowlists(), rows); + EXPECT_EQ(allowlists.dataset_rows(), dataset_rows); + EXPECT_EQ(allowlists.total_cardinality(), ids.size()); + EXPECT_EQ(allowlists.cardinality(0), 64); + EXPECT_TRUE(allowlists.empty(1)); + EXPECT_EQ(allowlists.cardinality(2), 128); + EXPECT_EQ(allowlists.cardinality(3), 64); + EXPECT_EQ(allowlists.view(1).device_reference(), nullptr); + EXPECT_NE(allowlists.view(0).device_reference(), allowlists.view(2).device_reference()); + EXPECT_EQ(device_allowlists.total_cardinality(), ids.size()); + EXPECT_EQ(device_allowlists.size_bytes(), allowlists.size_bytes()); + + expect_batch_membership(res, + allowlists, + 4, + {0, + 63, + 64, + 1000, + 0, + 1000, + dataset_rows - 1, + dataset_rows, + 999, + 1000, + 1254, + 1255, + 7, + (std::uint32_t{63} << 16) + 7, + 63, + dataset_rows}, + {1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0}); + expect_batch_membership(res, + device_allowlists, + 2, + {63, 64, 1, 1000, 1254, 1255, 7, (std::uint32_t{63} << 16) + 7}, + {1, 0, 0, 0, 1, 0, 1, 1}); + + std::array views{allowlists.view(0), allowlists.view(1), allowlists.view(2), allowlists.view(3)}; + cuvs::neighbors::filtering::roaring_filter filter(res, views); + EXPECT_EQ(filter.num_queries(), rows); + EXPECT_EQ(filter.cardinality(2), 128); +} + +TEST(RoaringAllowlist, BuildsGeneralMixedContainerRowsWithOneSegmentedSort) +{ + raft::device_resources res; + constexpr std::size_t rows = 3; + constexpr std::size_t width = 5000; + constexpr std::size_t dataset_rows = std::size_t{8} << 16; + std::vector const indptr{0, width, 2 * width, 3 * width}; + std::vector sorted; + sorted.reserve(rows * width); + for (std::uint32_t i = 0; i < width; ++i) { + sorted.push_back(100 + i); + } + for (std::uint32_t i = 0; i < width; ++i) { + sorted.push_back(2 * i); + } + for (std::uint32_t i = 0; i < width; ++i) { + auto const key = i % 8; + auto const value = (i / 8) * 2 + 1; + sorted.push_back((key << 16) + value); + } + for (std::size_t row = 0; row < rows; ++row) { + std::sort(sorted.begin() + indptr[row], sorted.begin() + indptr[row + 1]); + } + auto unsorted = sorted; + for (std::size_t row = 0; row < rows; ++row) { + std::reverse(unsorted.begin() + indptr[row], unsorted.begin() + indptr[row + 1]); + } + + auto general = from_ragged_ids(res, dataset_rows, unsorted, indptr, false); + auto presorted = from_ragged_ids(res, dataset_rows, sorted, indptr, true); + EXPECT_EQ(general.num_allowlists(), rows); + EXPECT_EQ(general.total_cardinality(), rows * width); + EXPECT_EQ(general.size_bytes(), presorted.size_bytes()); + expect_batch_membership(res, + general, + 5, + {99, + 100, + 5099, + 5100, + 9998, + 0, + 1, + 8192, + 9998, + 9999, + 1, + 2, + (std::uint32_t{7} << 16) + 1249, + (std::uint32_t{7} << 16) + 1250, + dataset_rows}, + {0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0}); +} + +TEST(RoaringAllowlist, ImportsRaggedPortableRowsIntoOnePackedOwner) +{ + raft::device_resources res; + auto array = array_row({1, 3, 5}); + auto run = run_row(100, 99); + std::vector bytes; + bytes.insert(bytes.end(), array.begin(), array.end()); + auto const empty_offset = bytes.size(); + bytes.insert(bytes.end(), run.begin(), run.end()); + std::array const offsets{0, array.size(), empty_offset, bytes.size()}; + auto allowlists = roaring_allowlist::from_serialized( + res, + 1000, + raft::make_host_vector_view(bytes.data(), bytes.size()), + raft::make_host_vector_view(offsets.data(), offsets.size())); + EXPECT_EQ(allowlists.num_allowlists(), 3); + EXPECT_EQ(allowlists.cardinality(0), 3); + EXPECT_TRUE(allowlists.empty(1)); + EXPECT_EQ(allowlists.cardinality(2), 100); + EXPECT_EQ(allowlists.view(1).device_reference(), nullptr); + expect_batch_membership(res, + allowlists, + 4, + {1, 2, 3, 5, 1, 3, 100, 999, 99, 100, 199, 200}, + {1, 0, 1, 1, 0, 0, 0, 0, 0, 1, 1, 0}); +} + +TEST(RoaringAllowlist, ConstructionPathsEmitByteIdenticalPortableRows) +{ + raft::device_resources res; + constexpr std::uint32_t second_key = std::uint32_t{1} << 16; + constexpr std::uint32_t third_key = std::uint32_t{2} << 16; + + std::vector ids; + for (std::uint32_t id = 100; id < 300; ++id) { + ids.push_back(id); // run + } + ids.insert(ids.end(), {second_key + 1, second_key + 3, second_key + 5}); // array + for (std::uint32_t id = 0; id < 10000; id += 2) { + ids.push_back(third_key + id); // bitmap + } + std::reverse(ids.begin(), ids.end()); + auto sorted = ids; + std::sort(sorted.begin(), sorted.end()); + auto const dataset_rows = static_cast(third_key) + 10000; + auto const expected = cuvs::bench::allowlist::make_portable_roaring_bytes(ids); + + auto host_unsorted = from_ids(res, dataset_rows, ids); + auto host_pre_sorted = from_ids(res, dataset_rows, sorted, true); + auto device_unsorted = from_device_ids(res, dataset_rows, ids); + auto device_sorted = from_device_ids(res, dataset_rows, sorted, true); + + EXPECT_EQ(copy_serialized_bytes(res, host_unsorted), expected); + EXPECT_EQ(copy_serialized_bytes(res, host_pre_sorted), expected); + EXPECT_EQ(copy_serialized_bytes(res, device_unsorted), expected); + EXPECT_EQ(copy_serialized_bytes(res, device_sorted), expected); +} + +TEST(RoaringAllowlist, BenchmarkPortableSerializerImportsMixedContainerKinds) +{ + raft::device_resources res; + constexpr std::uint32_t second_key = std::uint32_t{1} << 16; + constexpr std::uint32_t third_key = std::uint32_t{2} << 16; + + std::vector ids; + for (std::uint32_t id = 100; id < 300; ++id) { + ids.push_back(id); // run + } + ids.insert(ids.end(), {second_key + 1, second_key + 3, second_key + 5}); // array + for (std::uint32_t id = 0; id < 10000; id += 2) { + ids.push_back(third_key + id); // bitmap + } + std::reverse(ids.begin(), ids.end()); + auto const bytes = cuvs::bench::allowlist::make_portable_roaring_bytes(ids); + auto imported = import_row(res, static_cast(third_key) + 10000, bytes); + + EXPECT_EQ(imported.cardinality(0), ids.size()); + expect_membership(res, + imported, + {99, + 100, + 299, + 300, + second_key + 1, + second_key + 2, + third_key, + third_key + 1, + third_key + 9998, + third_key + 9999}, + {0, 1, 1, 0, 1, 0, 1, 0, 1, 0}); +} + +TEST(RoaringAllowlist, BuildsFromUnsortedUniqueDeviceIdsWithoutModifyingInput) +{ + raft::device_resources res; + constexpr std::uint32_t second_key = std::uint32_t{1} << 16; + constexpr std::uint32_t third_key = std::uint32_t{2} << 16; + + std::vector ids; + for (std::uint32_t id = 100; id < 300; ++id) { + ids.push_back(id); // run container + } + ids.insert(ids.end(), {second_key + 1, second_key + 3, second_key + 5}); // array container + for (std::uint32_t id = 0; id < 10000; id += 2) { + ids.push_back(third_key + id); // bitmap container + } + std::reverse(ids.begin(), ids.end()); + auto const original = ids; + + auto device_ids = raft::make_device_vector(res, ids.size()); + auto device_indptr = raft::make_device_vector(res, 2); + std::array const indptr{0, static_cast(ids.size())}; + auto const stream = raft::resource::get_cuda_stream(res); + raft::update_device(device_ids.data_handle(), ids.data(), ids.size(), stream); + raft::update_device(device_indptr.data_handle(), indptr.data(), indptr.size(), stream); + auto allowlist = + roaring_allowlist::from_ids(res, + third_key + 10000, + raft::make_device_vector_view( + device_ids.data_handle(), ids.size()), + raft::make_device_vector_view( + device_indptr.data_handle(), 2)); + + std::vector unchanged(ids.size()); + raft::update_host(unchanged.data(), device_ids.data_handle(), unchanged.size(), stream); + raft::resource::sync_stream(res); + EXPECT_EQ(unchanged, original); + EXPECT_EQ(allowlist.cardinality(0), 5203); + expect_membership(res, + allowlist, + {99, + 100, + 299, + 300, + second_key + 1, + second_key + 2, + third_key, + third_key + 1, + third_key + 9998, + third_key + 9999}, + {0, 1, 1, 0, 1, 0, 1, 0, 1, 0}); +} + +TEST(RoaringAllowlist, PreSortedFastPathBuildsHostAndDeviceIds) +{ + raft::device_resources res; + std::vector const sorted_ids{1, 2, 3, 65537, 65539, 131072}; + auto host_allowlist = from_ids(res, 131073, sorted_ids, true); + auto device_allowlist = from_device_ids(res, 131073, sorted_ids, true); + EXPECT_EQ(host_allowlist.cardinality(0), sorted_ids.size()); + EXPECT_EQ(device_allowlist.cardinality(0), sorted_ids.size()); + EXPECT_EQ(host_allowlist.size_bytes(), device_allowlist.size_bytes()); + expect_membership(res, host_allowlist, {0, 1, 3, 4, 65537, 65538, 131072}, {0, 1, 1, 0, 1, 0, 1}); + expect_membership( + res, device_allowlist, {0, 1, 3, 4, 65537, 65538, 131072}, {0, 1, 1, 0, 1, 0, 1}); + + auto empty = from_device_ids(res, 10, {}, true); + EXPECT_TRUE(empty.empty(0)); + EXPECT_THROW(from_device_ids(res, 10, {0, 9, 10}, true), raft::logic_error); +} + +TEST(RoaringAllowlist, DeviceFactoryHandlesEmptyAndRejectsOutOfRangeIds) +{ + raft::device_resources res; + auto empty = from_device_ids(res, 10, {}); + EXPECT_TRUE(empty.empty(0)); + EXPECT_EQ(empty.size_bytes(), sizeof(void const*) + sizeof(std::uint8_t)); + expect_membership(res, empty, {0, 9, 10}, {0, 0, 0}); + + EXPECT_THROW(from_device_ids(res, 10, {0, 9, 10}), raft::logic_error); + EXPECT_THROW(from_device_ids(res, 10, {0, 65536, 131072}), raft::logic_error); + std::vector general_invalid(129); + for (std::uint32_t i = 0; i < 128; ++i) { + general_invalid[i] = i; + } + general_invalid.back() = std::numeric_limits::max(); + EXPECT_THROW(from_device_ids(res, 1000, general_invalid), raft::logic_error); +} + +TEST(RoaringAllowlist, SparseDeviceBuilderHandlesThresholdAndGeneralCrossover) +{ + raft::device_resources res; + constexpr auto second_key = std::uint32_t{1} << 16; + + std::vector sorted; + sorted.reserve(129); + for (std::uint32_t id = 100; id < 164; ++id) { + sorted.push_back(id); // one run container + } + for (std::uint32_t value = 1; value < 65; value += 2) { + sorted.push_back(second_key + value); // one array container + } + for (std::uint32_t key = 2; sorted.size() < 129; ++key) { + sorted.push_back((key << 16) + 7); // many one-value array containers + } + ASSERT_EQ(sorted.size(), 129); + + std::vector threshold_ids(sorted.begin(), sorted.begin() + 128); + auto reversed_threshold = threshold_ids; + std::reverse(reversed_threshold.begin(), reversed_threshold.end()); + auto reversed_general = sorted; + std::reverse(reversed_general.begin(), reversed_general.end()); + auto const dataset_rows = static_cast(sorted.back()) + 1; + + auto sparse_unsorted = from_device_ids(res, dataset_rows, reversed_threshold); + auto general_unsorted = from_device_ids(res, dataset_rows, reversed_general); + std::vector pre_sorted_ids(sorted.begin(), sorted.begin() + 64); + auto sparse_pre_sorted = from_device_ids(res, dataset_rows, pre_sorted_ids, true); + auto general_pre_sorted = from_device_ids(res, dataset_rows, sorted, true); + + EXPECT_EQ(sparse_unsorted.cardinality(0), 128); + EXPECT_EQ(general_unsorted.cardinality(0), 129); + EXPECT_EQ(sparse_pre_sorted.cardinality(0), 64); + EXPECT_EQ(general_pre_sorted.cardinality(0), 129); + EXPECT_EQ(general_unsorted.size_bytes(), general_pre_sorted.size_bytes()); + + auto const sparse_last = threshold_ids.back(); + auto const general_last = sorted.back(); + expect_membership(res, + sparse_unsorted, + {99, 100, 163, 164, second_key + 1, second_key + 2, sparse_last, general_last}, + {0, 1, 1, 0, 1, 0, 1, 0}); + expect_membership(res, + general_unsorted, + {99, 100, 163, 164, second_key + 1, second_key + 2, sparse_last, general_last}, + {0, 1, 1, 0, 1, 0, 1, 1}); + expect_membership(res, sparse_pre_sorted, {99, 100, 163, 164}, {0, 1, 1, 0}); + expect_membership( + res, general_pre_sorted, {99, 100, 163, 164, second_key + 1, general_last}, {0, 1, 1, 0, 1, 1}); +} + +TEST(RoaringAllowlist, SparseDeviceSortPreservesMaximumUint32Id) +{ + raft::device_resources res; + auto const max_id = std::numeric_limits::max(); + auto allowlist = + from_device_ids(res, static_cast(std::uint64_t{1} << 32), {max_id, 0}); + EXPECT_EQ(allowlist.cardinality(0), 2); + expect_membership(res, allowlist, {0, 1, max_id - 1, max_id}, {1, 0, 0, 1}); +} + +TEST(RoaringAllowlist, ViewIsZeroCopyAndSurvivesOwnerMove) +{ + raft::device_resources res; + auto allowlist = from_ids(res, 32, {1, 4, 7}); + auto first_view = allowlist.view(0); + auto next_view = allowlist.view(0); + EXPECT_TRUE(first_view.valid()); + EXPECT_EQ(first_view.device_reference(), next_view.device_reference()); + EXPECT_NE(first_view.device_reference(), nullptr); + + auto moved = std::move(allowlist); + auto moved_view = moved.view(0); + 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}); + + auto empty = from_ids(res, 32, {}); + EXPECT_TRUE(empty.view(0).valid()); + EXPECT_EQ(empty.view(0).device_reference(), nullptr); + EXPECT_TRUE(empty.empty(0)); + EXPECT_EQ(empty.size_bytes(), sizeof(void const*) + sizeof(std::uint8_t)); + expect_membership(res, empty, {0, 31, 32}, {0, 0, 0}); +} + +TEST(RoaringAllowlist, StreamOrderedConstructionSupportsExplicitEventHandoff) +{ + 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, raft::resource::get_cuda_stream(build_res))); + RAFT_CUDA_TRY(cudaStreamWaitEvent(raft::resource::get_cuda_stream(consume_res), ready)); + expect_membership(consume_res, allowlist, {99, 100, 200, 900, 901}, {0, 1, 1, 1, 0}, true); + + // Serialized import has the same lifetime rule. Keep the caller-owned host bytes alive until the + // event recorded after construction has completed, then consume the initialized reference on the + // second stream. + auto bytes = array_row({2, 4, 8}); + auto imported = import_row(build_res, 1000, bytes); + RAFT_CUDA_TRY(cudaEventRecord(ready, raft::resource::get_cuda_stream(build_res))); + RAFT_CUDA_TRY(cudaStreamWaitEvent(raft::resource::get_cuda_stream(consume_res), ready)); + expect_membership(consume_res, imported, {1, 2, 4, 7, 8}, {0, 1, 1, 0, 1}, true); + + // Releasing an owner immediately is safe: device_uvector schedules the packed allocation's + // deallocation after encoding and reference initialization on the construction stream. + { + auto temporary = from_ids(build_res, 1000, {1, 10, 100}); + EXPECT_NE(temporary.view(0).device_reference(), nullptr); + } + RAFT_CUDA_TRY(cudaEventRecord(ready, raft::resource::get_cuda_stream(build_res))); + RAFT_CUDA_TRY(cudaEventSynchronize(ready)); + RAFT_CUDA_TRY(cudaEventDestroy(ready)); +} + +TEST(RoaringAllowlist, ImportsPortableArrayRunAndEmptyRows) +{ + raft::device_resources res; + + auto array_bytes = array_row({1, 3, 5}); + auto array = import_row(res, 1000, array_bytes); + EXPECT_EQ(array.cardinality(0), 3); + expect_membership(res, array, {1, 3, 4, 5}, {1, 1, 0, 1}); + + auto run_bytes = run_row(100, 99); + auto run = import_row(res, 1000, run_bytes); + EXPECT_EQ(run.cardinality(0), 100); + expect_membership(res, run, {99, 100, 199, 200}, {0, 1, 1, 0}, true); + + std::vector standard_empty; + append_u32(standard_empty, kCookieNoRun); + append_u32(standard_empty, 0); + auto empty = import_row(res, 1000, standard_empty); + EXPECT_TRUE(empty.empty(0)); + EXPECT_EQ(empty.view(0).device_reference(), nullptr); + + auto zero_length = import_row(res, 1000, {}); + EXPECT_TRUE(zero_length.empty(0)); +} + +TEST(RoaringAllowlist, RejectsSmallMalformedPortableInputs) +{ + raft::device_resources res; + + auto truncated = array_row({1, 3, 5}); + truncated.pop_back(); + EXPECT_THROW(import_row(res, 1000, truncated), raft::logic_error); + + auto bad_cookie = array_row({1}); + bad_cookie[0] = std::byte{0}; + EXPECT_THROW(import_row(res, 1000, bad_cookie), raft::logic_error); + + auto out_of_range = array_row({1000}); + EXPECT_THROW(import_row(res, 1000, out_of_range), raft::logic_error); + + auto valid = array_row({1, 3, 5}); + auto bytes = + raft::make_host_vector_view(valid.data(), valid.size()); + std::array const too_short{0}; + EXPECT_THROW(roaring_allowlist::from_serialized( + res, + 1000, + bytes, + raft::make_host_vector_view(too_short.data(), + too_short.size())), + raft::logic_error); + std::array const nonzero_start{1, valid.size()}; + EXPECT_THROW(roaring_allowlist::from_serialized( + res, + 1000, + bytes, + raft::make_host_vector_view( + nonzero_start.data(), nonzero_start.size())), + raft::logic_error); + std::array const decreasing{0, valid.size(), valid.size() - 1, valid.size()}; + EXPECT_THROW(roaring_allowlist::from_serialized( + res, + 1000, + bytes, + raft::make_host_vector_view(decreasing.data(), + decreasing.size())), + raft::logic_error); +} + +TEST(RoaringAllowlist, RejectsIdsOutsideLogicalShape) +{ + raft::device_resources res; + EXPECT_THROW(from_ids(res, 10, {10}), raft::logic_error); + EXPECT_THROW(from_ids(res, 0, {}), raft::logic_error); +} + +TEST(RoaringAllowlist, RejectsMalformedIndptr) +{ + raft::device_resources res; + std::vector const ids{1, 2}; + EXPECT_THROW(from_ragged_ids(res, 10, ids, {1, 2}), raft::logic_error); + EXPECT_THROW(from_ragged_ids(res, 10, ids, {0, 2, 1, 2}), raft::logic_error); + EXPECT_THROW(from_ragged_ids(res, 10, ids, {0, 1}), raft::logic_error); +} + +TEST(RoaringFilter, ReusesViewsAndUpdatesOneQueryOutsideSearch) +{ + 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(0), second.view(0), first.view(0)}; + 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); + EXPECT_EQ(filter.cardinality(2), 2); + EXPECT_FLOAT_EQ(filter.filtering_rate(), 0.875f); + EXPECT_GT(filter.size_bytes(), 0); + + auto const* payload = filter.device_payload(); + auto shared_copy = filter; + shared_copy.set_allowlist(res, 1, empty.view(0)); + EXPECT_EQ(filter.device_payload(), payload); + EXPECT_TRUE(filter.empty(1)); + EXPECT_FLOAT_EQ(filter.filtering_rate(), 0.999f); + + shared_copy.set_allowlist(res, 1, second.view(0)); + EXPECT_EQ(filter.cardinality(1), 3); + EXPECT_EQ(filter.device_payload(), payload); +} + +TEST(RoaringFilter, RejectsSmallInvalidMappings) +{ + raft::device_resources res; + EXPECT_THROW(cuvs::neighbors::filtering::roaring_filter( + res, std::span{}), + raft::logic_error); + + auto ten = from_ids(res, 10, {1}); + auto eleven = from_ids(res, 11, {1}); + std::array mismatched{ten.view(0), eleven.view(0)}; + EXPECT_THROW(cuvs::neighbors::filtering::roaring_filter(res, mismatched), raft::logic_error); + + std::array invalid{cuvs::core::roaring_allowlist_view{}}; + EXPECT_THROW(cuvs::neighbors::filtering::roaring_filter(res, invalid), raft::logic_error); + + std::array one{ten.view(0)}; + cuvs::neighbors::filtering::roaring_filter filter(res, one); + EXPECT_THROW(filter.set_allowlist(res, 1, ten.view(0)), raft::logic_error); + EXPECT_THROW(filter.set_allowlist(res, 0, eleven.view(0)), raft::logic_error); +} + +TEST(CagraFilterBenchmarkUtils, GeneratesExactDeterministicAllowlistLayouts) +{ + using cuvs::bench::cagra_filter::allowlist_layout; + using cuvs::bench::cagra_filter::make_allowed_ids; + + for (auto layout : {allowlist_layout::random, allowlist_layout::runs}) { + auto first = make_allowed_ids(257, 73, layout, 12345); + auto second = make_allowed_ids(257, 73, layout, 12345); + EXPECT_EQ(first, second); + ASSERT_EQ(first.size(), 73); + + std::sort(first.begin(), first.end()); + EXPECT_EQ(std::adjacent_find(first.begin(), first.end()), first.end()); + EXPECT_LT(first.back(), 257); + } + + auto run = make_allowed_ids(257, 73, allowlist_layout::runs, 12345); + std::sort(run.begin(), run.end()); + auto breaks = std::size_t{0}; + for (std::size_t i = 1; i < run.size(); ++i) { + breaks += run[i] != run[i - 1] + 1 ? 1 : 0; + } + EXPECT_LE(breaks, + 1); // a circular interval is one range, or two when it wraps at row zero + + auto random = make_allowed_ids(257, 73, allowlist_layout::random, 12345); + std::sort(random.begin(), random.end()); + breaks = 0; + for (std::size_t i = 1; i < random.size(); ++i) { + breaks += random[i] != random[i - 1] + 1 ? 1 : 0; + } + EXPECT_GT(breaks, 1); + + EXPECT_TRUE(make_allowed_ids(257, 0, allowlist_layout::random, 12345).empty()); + EXPECT_THROW(make_allowed_ids(0, 0, allowlist_layout::random, 12345), std::invalid_argument); + EXPECT_THROW(make_allowed_ids(10, 11, allowlist_layout::runs, 12345), std::invalid_argument); +} + +TEST(CagraFilterBenchmarkUtils, SummarizesLatencySamples) +{ + auto const stats = cuvs::bench::cagra_filter::summarize({1.0, 2.0, 3.0, 4.0, 5.0}); + EXPECT_DOUBLE_EQ(stats.mean, 3.0); + EXPECT_DOUBLE_EQ(stats.p50, 3.0); + EXPECT_DOUBLE_EQ(stats.p95, 4.8); + EXPECT_DOUBLE_EQ(stats.p99, 4.96); + + EXPECT_EQ(cuvs::bench::cagra_filter::cardinality_for_fraction(1'000'000, 0.001, 10), 1000); + EXPECT_EQ(cuvs::bench::cagra_filter::allowed_percent_label(0.001), "allowed_0.1pct"); + EXPECT_EQ(cuvs::bench::cagra_filter::allowed_percent_label(0.015), "allowed_1.5pct"); +} + +TEST(AllowlistBenchmarkUtils, GeneratesControlledSmallLookupProfiles) +{ + using namespace cuvs::bench::allowlist; + + constexpr std::uint32_t rows = 64; + constexpr std::size_t cardinality = 16; + for (auto profile : {lookup_profile::best, lookup_profile::average, lookup_profile::worst}) { + auto first = make_allowed_ids(rows, cardinality, profile, 12345); + auto second = make_allowed_ids(rows, cardinality, profile, 12345); + EXPECT_EQ(first, second); + ASSERT_EQ(first.size(), cardinality); + std::sort(first.begin(), first.end()); + EXPECT_EQ(std::adjacent_find(first.begin(), first.end()), first.end()); + EXPECT_LT(first.back(), rows); + } + + auto best = make_allowed_ids(rows, cardinality, lookup_profile::best, 12345); + EXPECT_EQ(occupied_chunks(best), 1); + EXPECT_EQ(count_runs(best), 1); + + auto worst = make_allowed_ids(rows, cardinality, lookup_profile::worst, 12345); + EXPECT_EQ(maximum_encoded_runs(cardinality, rows), 7); + EXPECT_EQ(count_runs(worst), 7); + + EXPECT_THROW(make_allowed_ids(0, 0, lookup_profile::best, 12345), std::invalid_argument); + EXPECT_THROW(make_allowed_ids(10, 11, lookup_profile::worst, 12345), std::invalid_argument); +} + +TEST(AllowlistBenchmarkUtils, GeneratesBestHitsAverageDomainAndWorstMisses) +{ + using namespace cuvs::bench::allowlist; + constexpr std::uint32_t rows = 64; + constexpr std::size_t cardinality = 16; + for (auto profile : {lookup_profile::best, lookup_profile::average, lookup_profile::worst}) { + auto ids = make_allowed_ids(rows, cardinality, profile, 12345); + auto probes = make_probes(rows, ids, profile, 257, 67890); + ASSERT_EQ(probes.size(), 257); + EXPECT_TRUE(std::all_of(probes.begin(), probes.end(), [](auto probe) { return probe < rows; })); + std::sort(ids.begin(), ids.end()); + if (profile == lookup_profile::best) { + EXPECT_TRUE(std::all_of(probes.begin(), probes.end(), [&](auto probe) { + return std::binary_search(ids.begin(), ids.end(), probe); + })); + } else if (profile == lookup_profile::worst) { + EXPECT_TRUE(std::none_of(probes.begin(), probes.end(), [&](auto probe) { + return std::binary_search(ids.begin(), ids.end(), probe); + })); + } + } + EXPECT_THROW(make_probes(rows, {1}, lookup_profile::average, 0, 67890), std::invalid_argument); + + auto full_domain = make_allowed_ids(std::uint64_t{1} << 32, 1024, lookup_profile::average, 12345); + auto repeated = make_allowed_ids(std::uint64_t{1} << 32, 1024, lookup_profile::average, 12345); + EXPECT_EQ(full_domain, repeated); + std::sort(full_domain.begin(), full_domain.end()); + EXPECT_EQ(std::adjacent_find(full_domain.begin(), full_domain.end()), full_domain.end()); +} + +} // 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..1a65d940c6 --- /dev/null +++ b/examples/cpp/src/cagra_roaring_filter_example.cu @@ -0,0 +1,174 @@ +/* + * 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; +constexpr std::uint32_t portable_cookie_no_run = 12346; +constexpr std::uint32_t single_array_payload_byte_offset = 16; + +void append_u16(std::vector& bytes, std::uint16_t value) +{ + bytes.push_back(static_cast(value & 0xffu)); + bytes.push_back(static_cast((value >> 8) & 0xffu)); +} + +void append_u32(std::vector& bytes, std::uint32_t value) +{ + for (int i = 0; i < 4; ++i) { + bytes.push_back(static_cast((value >> (8 * i)) & 0xffu)); + } +} + +// This helper intentionally writes the simplest nonempty portable Roaring row: +// +// bytes 0..3: uint32 no-run cookie (12346) +// bytes 4..7: uint32 container count (1) +// bytes 8..11: { uint16 key, uint16 cardinality_minus_one } +// bytes 12..15: uint32 payload offset, measured from byte 0 of this row +// bytes 16..: sorted uint16 array values +// +// Every integer is little-endian. This example's IDs are below 2^16, so they all use container key +// zero; a general encoder must split IDs by their high 16 bits and may need bitmap or run +// containers. Prefer `from_ids` when starting with IDs. `from_serialized` is intended primarily +// for interoperating with systems that already emit the standard format: +// +// https://github.com/RoaringBitmap/RoaringFormatSpec +void append_portable_array_row(std::vector& bytes, + std::vector const& ids, + std::int64_t first, + std::int64_t last) +{ + append_u32(bytes, portable_cookie_no_run); + append_u32(bytes, 1); // one container + append_u16(bytes, 0); // high 16-bit container key + append_u16(bytes, static_cast(last - first - 1)); + append_u32(bytes, single_array_payload_byte_offset); + for (auto i = first; i < last; ++i) { + append_u16(bytes, static_cast(ids[static_cast(i)])); + } +} + +} // 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); + index.update_device_dataset_same_layout(res, padded); + + // Ragged construction: query q accepts rows whose row id modulo n_queries equals q. + std::vector ids; + std::vector indptr{0}; + for (std::uint32_t query = 0; query < n_queries; ++query) { + for (std::uint32_t row = query; row < n_rows; row += n_queries) { + ids.push_back(row); + } + indptr.push_back(static_cast(ids.size())); + } + // One factory call consumes the ragged rows directly and retains every variable-length stream + // in one owner. + auto id_allowlists = cuvs::core::roaring_allowlist::from_ids( + res, + n_rows, + raft::make_host_vector_view(ids.data(), ids.size()), + raft::make_host_vector_view(indptr.data(), indptr.size()), + true); + std::vector id_views; + for (std::int64_t query = 0; query < n_queries; ++query) { + id_views.push_back(id_allowlists.view(query)); + } + cuvs::neighbors::filtering::roaring_filter ids_filter(res, id_views); + + // A serialized system can keep ragged portable rows in one host buffer plus outer offsets. The + // whole packed buffer becomes one owner even though encoded row sizes may differ. + std::vector bytes; + std::vector byte_offsets{0}; + for (std::int64_t query = 0; query < n_queries; ++query) { + append_portable_array_row(bytes, ids, indptr[query], indptr[query + 1]); + byte_offsets.push_back(bytes.size()); + } + + auto serialized_allowlists = cuvs::core::roaring_allowlist::from_serialized( + res, + n_rows, + raft::make_host_vector_view(bytes.data(), bytes.size()), + raft::make_host_vector_view(byte_offsets.data(), + byte_offsets.size())); + std::vector serialized_views; + for (std::int64_t query = 0; query < n_queries; ++query) { + serialized_views.push_back(serialized_allowlists.view(query)); + } + cuvs::neighbors::filtering::roaring_filter filter(res, serialized_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 << "Built " << id_allowlists.num_allowlists() << " ID allowlists and imported " + << serialized_allowlists.num_allowlists() + << " portable allowlists; CAGRA reused one prepared filter payload.\n"; + return ids_filter.num_queries() == n_queries ? 0 : 1; +} From 35f50f523d2035d24278d96b2337d90e6b0f4d36 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Wed, 2 Sep 2026 19:49:11 +0000 Subject: [PATCH 2/5] single batch upstream cuco --- cpp/include/cuvs/core/roaring_allowlist.hpp | 18 +- cpp/src/core/roaring_allowlist.cu | 682 +++++++------------- cpp/tests/neighbors/roaring_allowlist.cu | 213 ++---- 3 files changed, 302 insertions(+), 611 deletions(-) diff --git a/cpp/include/cuvs/core/roaring_allowlist.hpp b/cpp/include/cuvs/core/roaring_allowlist.hpp index af190454f9..60789d5956 100644 --- a/cpp/include/cuvs/core/roaring_allowlist.hpp +++ b/cpp/include/cuvs/core/roaring_allowlist.hpp @@ -62,13 +62,18 @@ class CUVS_EXPORT roaring_allowlist_view { * * Logically, the owner is a sparse matrix with one allowlist row per query and one possible column * per dataset row. @ref from_ids accepts one contiguous ID vector plus an indptr vector that - * delimits independently sized query rows. Every row is sorted and encoded independently. All - * variable-length portable Roaring streams and their initialized + * delimits independently sized query rows. Every row is sorted and encoded independently. For + * multiple rows, all variable-length portable Roaring streams and their initialized * `cuco::experimental::roaring_bitmap_ref` objects share one packed device allocation. + * The batch builder uses indptr directly for segmented device radix sort and schedules + * analysis/encoding over all containers in all rows. * - * For multiple rows, the builder uses indptr directly for segmented device radix sort and - * schedules analysis/encoding over all containers in all rows. A one-row input retains the tuned - * single-allowlist builder. Final encoding and reference initialization remain stream ordered. + * A one-row input delegates raw-index construction to cuCollections. cuVS retains the cuco owner + * directly and materializes only its lightweight reference in cuVS device metadata; the serialized + * payload is not copied. cuco currently emits array and bitmap containers on this path. The cuVS + * multi-row builder emits the same array and bitmap container forms. Imported portable rows may + * still contain run containers. Final encoding and reference initialization remain stream ordered + * in both paths. * * IDs must be unique within each row. Setting @p pre_sorted skips sorting and promises that every * row is strictly increasing; ordering and uniqueness are not checked. Every ID must be smaller @@ -76,7 +81,8 @@ class CUVS_EXPORT roaring_allowlist_view { * * @see https://github.com/RoaringBitmap/RoaringFormatSpec * @see - * https://github.com/NVIDIA/cuCollections/blob/6001618aaa7f17ea2bbcd444650e9573c4f3d6c5/include/cuco/roaring_bitmap_ref.cuh + * https://github.com/NVIDIA/cuCollections/blob/9d7c9307395c3b8795d93ad65d0751c98471dde6/include/cuco/roaring_bitmap_ref.cuh + * @see https://github.com/NVIDIA/cuCollections/pull/839 */ class CUVS_EXPORT roaring_allowlist { private: diff --git a/cpp/src/core/roaring_allowlist.cu b/cpp/src/core/roaring_allowlist.cu index 42a2778020..141f561e87 100644 --- a/cpp/src/core/roaring_allowlist.cu +++ b/cpp/src/core/roaring_allowlist.cu @@ -9,9 +9,18 @@ #include +// cuCollections PR #839 adds GPU construction from raw indices. Keep the +// existing cuVS builder as a compatibility fallback until that API reaches the +// pinned cuco revision, and keep using it for the segmented multi-row path. +#if __has_include() +#include +#include +#define CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER 1 +#else +#define CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER 0 +#endif + #include -#include -#include #include #include #include @@ -78,7 +87,7 @@ namespace { * container payloads * @endcode * - * With at least one run container, the row is laid out as: + * Imported rows with at least one run container are laid out as: * * @code{.unparsed} * uint32 cookie = 12347 | ((N - 1) << 16) @@ -111,6 +120,9 @@ namespace { * { uint16 start; uint16 length_minus_one; } runs[number_of_runs] * @endcode * + * ID-based construction emits only the array and bitmap forms above. The run + * form is accepted only when importing an existing standard-portable stream. + * * The portable format above describes exactly one bitmap. Query-to-allowlist * association is a separate concern: `filtering::roaring_filter` stores device * pointers to already initialized allowlist references. Consequently neither @@ -126,6 +138,11 @@ constexpr std::size_t kOffsetThreshold = 4; using ref_type = cuco::experimental::roaring_bitmap_ref; +#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER +using cuco_bitmap_allocator = rmm::mr::polymorphic_allocator; +using cuco_bitmap_type = cuco::experimental::roaring_bitmap; +#endif + struct row_metadata { std::size_t cardinality{}; std::uint32_t max_id{}; @@ -159,7 +176,7 @@ void validate_dataset_rows(std::size_t dataset_rows) "dataset_rows exceeds the uint32_t Roaring key domain."); } -enum class container_kind : std::uint8_t { array, bitmap, run }; +enum class container_kind : std::uint8_t { array, bitmap }; std::size_t align_up(std::size_t offset, std::size_t alignment) { @@ -170,7 +187,6 @@ struct device_build_summary { std::int64_t cardinality{}; std::uint64_t payload_bytes{}; std::uint32_t num_containers{}; - std::uint32_t has_run{}; std::uint32_t invalid{}; }; @@ -179,6 +195,9 @@ struct device_build_result { std::size_t serialized_bytes{}; std::size_t cardinality{}; bool reference_initialized{}; +#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER + std::unique_ptr cuco_owner{}; +#endif }; std::size_t reference_offset(std::size_t serialized_bytes) @@ -195,12 +214,11 @@ constexpr int kBuilderBlockSize = 256; // Small allowlists do not benefit from the general builder's device-wide sort, // two scans, and separate per-stage allocations. At this cardinality every -// portable container is necessarily an array or a run (a bitmap requires more -// than 4096 values in one high-16-bit partition), so one CTA can sort, analyze, -// and later encode the complete row. For a single pre-sorted row the cutoff is lower because the -// general path already avoids its most expensive stage, the device-wide sort. Batched rows use the -// 128-ID capacity because the launch is amortized across the matrix; keep both rules tied to the -// construction benchmark. +// portable container is necessarily an array (a bitmap requires more than 4096 values in one +// high-16-bit partition), so one CTA can sort, analyze, and later encode the complete row. For a +// single pre-sorted row the cutoff is lower because the general path already avoids its most +// expensive stage, the device-wide sort. Batched rows use the 128-ID capacity because the launch is +// amortized across the matrix; keep both rules tied to the construction benchmark. constexpr int kSparseBuilderBlockSize = 128; constexpr std::size_t kSparseBuilderMaxIds = 128; constexpr std::size_t kSparseBuilderMaxPreSortedIds = 64; @@ -211,12 +229,9 @@ static_assert(kSparseBuilderMaxIds % kSparseBuilderBlockSize == 0); struct sparse_container_metadata { std::uint32_t begin{}; std::uint32_t payload_offset{}; - std::uint16_t runs{}; - container_kind kind{}; - std::uint8_t padding{}; }; -static_assert(sizeof(sparse_container_metadata) == 12); +static_assert(sizeof(sparse_container_metadata) == 8); struct sparse_scratch_layout { explicit sparse_scratch_layout(std::size_t ids, std::size_t containers, bool store_sorted_ids) @@ -262,7 +277,6 @@ struct general_scratch_layout { kinds_offset = reserve(containers, sizeof(container_kind), alignof(container_kind)); payload_sizes_offset = reserve(containers, sizeof(std::uint64_t), alignof(std::uint64_t)); payload_offsets_offset = reserve(containers, sizeof(std::uint64_t), alignof(std::uint64_t)); - has_run_offset = reserve(1, sizeof(std::uint32_t), alignof(std::uint32_t)); summary_offset = reserve(1, sizeof(device_build_summary), alignof(device_build_summary)); workspace_offset = reserve(workspace_bytes, sizeof(cuda::std::byte), alignof(std::max_align_t)); bytes = cursor; @@ -277,7 +291,6 @@ struct general_scratch_layout { std::size_t kinds_offset{}; std::size_t payload_sizes_offset{}; std::size_t payload_offsets_offset{}; - std::size_t has_run_offset{}; std::size_t summary_offset{}; std::size_t workspace_offset{}; std::size_t bytes{}; @@ -353,55 +366,36 @@ __global__ void narrow_container_count_kernel(std::int64_t const* selected_count } /** - * Count consecutive runs and select the smallest legal portable payload for - * each container. + * Select the standard array or bitmap portable payload for each container. * - * Each block owns one container. Threads independently identify run starts in - * the sorted slice, then a block reduction produces the exact run count. This - * uses O(number of input IDs) scratch for sorting and scans; it never - * constructs a dense dataset-sized bitmap. + * ID construction intentionally does not emit run containers. Full and nearly + * full allowlists should normally bypass filtering, and limiting construction + * to the two cuco-native forms keeps the batch builder and lookup behavior + * predictable. This still uses O(number of input IDs) scratch for sorting and + * scans; it never constructs a dense dataset-sized temporary bitmap. */ -__global__ void analyze_containers_kernel(std::uint32_t const* ids, - std::int64_t const* id_count, +__global__ void analyze_containers_kernel(std::int64_t const* id_count, std::int64_t const* container_starts, std::uint32_t const* num_containers, container_kind* kinds, - std::uint64_t* payload_sizes, - std::uint32_t* has_run) + std::uint64_t* payload_sizes) { - auto const container = static_cast(blockIdx.x); - auto const count = *num_containers; - if (container >= count) { return; } - - auto const begin = container_starts[container]; - auto const end = container + 1 < count ? container_starts[container + 1] : *id_count; - std::uint32_t local_runs{}; - for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { - local_runs += - i == begin || static_cast(ids[i]) != static_cast(ids[i - 1]) + 1 - ? 1u - : 0u; - } - - using block_reduce = cub::BlockReduce; - __shared__ typename block_reduce::TempStorage reduction_storage; - auto const runs = block_reduce(reduction_storage).Sum(local_runs); - if (threadIdx.x != 0) { return; } - - auto const cardinality = static_cast(end - begin); - auto const normal_size = - cardinality <= kArrayCardinality ? cardinality * sizeof(std::uint16_t) : kBitmapBytes; - auto const run_size = sizeof(std::uint16_t) + runs * 2 * sizeof(std::uint16_t); - if (run_size < normal_size) { - kinds[container] = container_kind::run; - payload_sizes[container] = run_size; - atomicExch(has_run, 1u); - } else if (cardinality <= kArrayCardinality) { - kinds[container] = container_kind::array; - payload_sizes[container] = normal_size; - } else { - kinds[container] = container_kind::bitmap; - payload_sizes[container] = normal_size; + auto container = + static_cast(static_cast(blockIdx.x) * blockDim.x + threadIdx.x); + auto const stride = static_cast(static_cast(gridDim.x) * + static_cast(blockDim.x)); + auto const count = *num_containers; + for (; container < count; container += stride) { + auto const begin = container_starts[container]; + auto const end = container + 1 < count ? container_starts[container + 1] : *id_count; + auto const cardinality = static_cast(end - begin); + if (cardinality <= kArrayCardinality) { + kinds[container] = container_kind::array; + payload_sizes[container] = cardinality * sizeof(std::uint16_t); + } else { + kinds[container] = container_kind::bitmap; + payload_sizes[container] = kBitmapBytes; + } } } @@ -410,7 +404,6 @@ __global__ void analyze_containers_kernel(std::uint32_t const* ids, __global__ void finish_device_analysis_kernel(std::int64_t const* id_count, std::int64_t const* valid_count, std::uint32_t const* num_containers, - std::uint32_t const* has_run, std::uint64_t const* payload_sizes, std::uint64_t const* payload_offsets, device_build_summary* summary) @@ -420,64 +413,34 @@ __global__ void finish_device_analysis_kernel(std::int64_t const* id_count, auto const containers = *num_containers; summary->cardinality = cardinality; summary->num_containers = containers; - summary->has_run = *has_run; summary->payload_bytes = containers == 0 ? 0 : payload_offsets[containers - 1] + payload_sizes[containers - 1]; summary->invalid = cardinality != *valid_count; } -__host__ __device__ std::size_t portable_header_size(std::uint32_t num_containers, bool has_run) +__host__ __device__ std::size_t portable_header_size(std::uint32_t num_containers) { - if (!has_run) { - return 2 * sizeof(std::uint32_t) + - num_containers * (2 * sizeof(std::uint16_t) + sizeof(std::uint32_t)); - } - auto const run_bitmap_bytes = (num_containers + 7) / 8; - return sizeof(std::uint32_t) + run_bitmap_bytes + num_containers * 2 * sizeof(std::uint16_t) + - (num_containers >= kOffsetThreshold ? num_containers * sizeof(std::uint32_t) : 0); + return 2 * sizeof(std::uint32_t) + + num_containers * (2 * sizeof(std::uint16_t) + sizeof(std::uint32_t)); } -/** Write the cookie, run bitmap, descriptors, and portable container-offset - * table. */ +/** Write the cookie, descriptors, and portable container-offset table. */ __global__ void encode_header_kernel(std::uint32_t const* ids, std::int64_t const* id_count, std::int64_t const* container_starts, std::uint32_t num_containers, - container_kind const* kinds, std::uint64_t const* payload_offsets, std::size_t header_size, - bool has_run, cuda::std::byte* output) { - auto const thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - auto const stride = static_cast(gridDim.x) * blockDim.x; - auto const run_bitmap_bytes = has_run ? (num_containers + 7) / 8 : 0; - auto const descriptor_offset = - has_run ? sizeof(std::uint32_t) + run_bitmap_bytes : 2 * sizeof(std::uint32_t); - auto const offsets_offset = descriptor_offset + num_containers * 2 * sizeof(std::uint16_t); - bool const store_offsets = !has_run || num_containers >= kOffsetThreshold; + auto const thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + auto const stride = static_cast(gridDim.x) * blockDim.x; + constexpr auto descriptor_offset = 2 * sizeof(std::uint32_t); + auto const offsets_offset = descriptor_offset + num_containers * 2 * sizeof(std::uint16_t); if (thread == 0) { - if (has_run) { - write_u32(output, 0, kCookieRun | ((num_containers - 1) << 16)); - } else { - write_u32(output, 0, kCookieNoRun); - write_u32(output, sizeof(std::uint32_t), num_containers); - } - } - - // Write every run-bitmap byte directly. This initializes unused high bits to zero and avoids a - // memset of the complete serialized allocation. - auto* output_bytes = reinterpret_cast(output); - for (auto byte = thread; byte < run_bitmap_bytes; byte += stride) { - std::uint8_t value{}; - for (std::uint32_t bit = 0; bit < 8; ++bit) { - auto const container = static_cast(byte * 8 + bit); - if (container < num_containers && kinds[container] == container_kind::run) { - value |= static_cast(1u << bit); - } - } - output_bytes[sizeof(std::uint32_t) + byte] = value; + write_u32(output, 0, kCookieNoRun); + write_u32(output, sizeof(std::uint32_t), num_containers); } for (auto container = thread; container < num_containers; container += stride) { @@ -487,21 +450,18 @@ __global__ void encode_header_kernel(std::uint32_t const* ids, write_u16(output, descriptor, static_cast(ids[begin] >> 16)); write_u16( output, descriptor + sizeof(std::uint16_t), static_cast(end - begin - 1)); - if (store_offsets) { - write_u32(output, - offsets_offset + container * sizeof(std::uint32_t), - static_cast(header_size + payload_offsets[container])); - } + write_u32(output, + offsets_offset + container * sizeof(std::uint32_t), + static_cast(header_size + payload_offsets[container])); } } -/** Encode array, bitmap, and run payloads directly into their final device offsets. */ +/** Encode array and bitmap payloads directly into their final device offsets. */ __global__ void encode_payloads_kernel(std::uint32_t const* ids, std::int64_t const* id_count, std::int64_t const* container_starts, std::uint32_t num_containers, container_kind const* kinds, - std::uint64_t const* payload_sizes, std::uint64_t const* payload_offsets, std::size_t header_size, cuda::std::byte* output) @@ -521,75 +481,20 @@ __global__ void encode_payloads_kernel(std::uint32_t const* ids, return; } - using run_scan = cub::BlockScan; - union payload_scratch { - std::uint64_t bitmap_words[kBitmapBytes / sizeof(std::uint64_t)]; - typename run_scan::TempStorage run_scan_storage; - }; - __shared__ payload_scratch scratch; - __shared__ std::uint32_t run_base; - __shared__ std::uint32_t tile_runs; - - if (kinds[container] == container_kind::bitmap) { - constexpr std::uint32_t words = kBitmapBytes / sizeof(std::uint64_t); - for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { - scratch.bitmap_words[word] = 0; - } - __syncthreads(); - for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { - auto const lower = ids[i] & 0xffffu; - atomicOr(reinterpret_cast(&scratch.bitmap_words[lower / 64]), - static_cast(std::uint64_t{1} << (lower % 64))); - } - __syncthreads(); - for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { - write_u64(output, payload + word * sizeof(std::uint64_t), scratch.bitmap_words[word]); - } - return; + __shared__ std::uint64_t bitmap_words[kBitmapBytes / sizeof(std::uint64_t)]; + constexpr std::uint32_t words = kBitmapBytes / sizeof(std::uint64_t); + for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { + bitmap_words[word] = 0; } - - auto const num_runs = static_cast( - (payload_sizes[container] - sizeof(std::uint16_t)) / (2 * sizeof(std::uint16_t))); - if (threadIdx.x == 0) { - write_u16(output, payload, num_runs); - run_base = 0; + __syncthreads(); + for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { + auto const lower = ids[i] & 0xffffu; + atomicOr(reinterpret_cast(&bitmap_words[lower / 64]), + static_cast(std::uint64_t{1} << (lower % 64))); } __syncthreads(); - - // A block scan assigns stable output positions to run starts in each tile. Each run-start thread - // walks only its own run, so total work remains O(container cardinality) while high-run-count - // containers use the complete CTA instead of one serial thread. - for (auto tile = begin; tile < end; tile += blockDim.x) { - auto const i = tile + threadIdx.x; - std::uint32_t const is_run_start = - i < end && (i == begin || - static_cast(ids[i]) != static_cast(ids[i - 1]) + 1) - ? 1u - : 0u; - std::uint32_t run_rank{}; - std::uint32_t block_runs{}; - run_scan(scratch.run_scan_storage).ExclusiveSum(is_run_start, run_rank, block_runs); - if (threadIdx.x == 0) { tile_runs = block_runs; } - __syncthreads(); - - if (is_run_start != 0) { - auto j = i + 1; - while (j < end && - static_cast(ids[j]) == static_cast(ids[j - 1]) + 1) { - ++j; - } - auto const start = static_cast(ids[i] & 0xffffu); - auto const last = static_cast(ids[j - 1] & 0xffffu); - auto const run_offset = - payload + sizeof(std::uint16_t) + - static_cast(run_base + run_rank) * 2 * sizeof(std::uint16_t); - write_u16(output, run_offset, start); - write_u16( - output, run_offset + sizeof(std::uint16_t), static_cast(last - start)); - } - __syncthreads(); - if (threadIdx.x == 0) { run_base += tile_runs; } - __syncthreads(); + for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { + write_u64(output, payload + word * sizeof(std::uint64_t), bitmap_words[word]); } } @@ -639,7 +544,6 @@ __global__ void analyze_sparse_ids_kernel(std::uint32_t const* ids, summary->cardinality = size; summary->payload_bytes = 0; summary->num_containers = 0; - summary->has_run = 0; summary->invalid = 0; // Validate before writing container metadata. For a valid row, the logical @@ -657,27 +561,14 @@ __global__ void analyze_sparse_ids_kernel(std::uint32_t const* ids, while (begin < size) { auto const key = normalized_ids[begin] >> 16; auto end = begin + 1; - std::uint32_t runs{1}; while (end < size && (normalized_ids[end] >> 16) == key) { - runs += static_cast(normalized_ids[end]) != - static_cast(normalized_ids[end - 1]) + 1 - ? 1u - : 0u; ++end; } auto const cardinality = end - begin; auto const array_size = cardinality * sizeof(std::uint16_t); - auto const run_size = sizeof(std::uint16_t) + runs * 2 * sizeof(std::uint16_t); - auto const use_run = run_size < array_size; - metadata[container] = - sparse_container_metadata{begin, - payload_offset, - static_cast(runs), - use_run ? container_kind::run : container_kind::array, - 0}; - payload_offset += use_run ? run_size : array_size; - summary->has_run |= use_run ? 1u : 0u; + metadata[container] = sparse_container_metadata{begin, payload_offset}; + payload_offset += array_size; ++container; begin = end; } @@ -687,50 +578,26 @@ __global__ void analyze_sparse_ids_kernel(std::uint32_t const* ids, } /** - * Encode a sparse row in one CTA after the host has allocated the exact byte - * count reported by `analyze_sparse_ids_kernel`. + * Encode a sparse row in one CTA after exact output allocation. * - * Thread zero writes every header byte, including the run bitmap, so this path - * needs no output memset. Array values are striped across the CTA; thread zero - * writes the comparatively small run payloads. Bitmap payloads cannot occur - * below the sparse cardinality threshold. + * Thread zero writes every header byte, so this path needs no output memset. + * Array values are striped across the CTA. Bitmap payloads cannot occur below + * the sparse cardinality threshold. */ __global__ void encode_sparse_row_kernel(std::uint32_t const* ids, std::uint32_t cardinality, sparse_container_metadata const* metadata, std::uint32_t num_containers, std::size_t header_size, - bool has_run, cuda::std::byte* output, ref_type* reference) { - bool const store_offsets = !has_run || num_containers >= kOffsetThreshold; - std::size_t descriptor_offset{}; - std::size_t offsets_offset{}; + constexpr auto descriptor_offset = 2 * sizeof(std::uint32_t); + auto const offsets_offset = descriptor_offset + num_containers * 2 * sizeof(std::uint16_t); if (threadIdx.x == 0) { - if (has_run) { - write_u32(output, 0, kCookieRun | ((num_containers - 1) << 16)); - auto const run_bitmap_bytes = (num_containers + 7) / 8; - for (std::uint32_t byte = 0; byte < run_bitmap_bytes; ++byte) { - std::uint8_t value{}; - for (std::uint32_t bit = 0; bit < 8; ++bit) { - auto const container = byte * 8 + bit; - if (container < num_containers && metadata[container].kind == container_kind::run) { - value |= static_cast(1u << bit); - } - } - reinterpret_cast(output)[sizeof(std::uint32_t) + byte] = value; - } - descriptor_offset = sizeof(std::uint32_t) + run_bitmap_bytes; - offsets_offset = descriptor_offset + num_containers * 2 * sizeof(std::uint16_t); - } else { - write_u32(output, 0, kCookieNoRun); - write_u32(output, sizeof(std::uint32_t), num_containers); - descriptor_offset = 2 * sizeof(std::uint32_t); - offsets_offset = descriptor_offset + num_containers * 2 * sizeof(std::uint16_t); - } - + write_u32(output, 0, kCookieNoRun); + write_u32(output, sizeof(std::uint32_t), num_containers); for (std::uint32_t container = 0; container < num_containers; ++container) { auto const begin = metadata[container].begin; auto const end = container + 1 < num_containers ? metadata[container + 1].begin : cardinality; @@ -738,11 +605,9 @@ __global__ void encode_sparse_row_kernel(std::uint32_t const* ids, write_u16(output, descriptor, static_cast(ids[begin] >> 16)); write_u16( output, descriptor + sizeof(std::uint16_t), static_cast(end - begin - 1)); - if (store_offsets) { - write_u32(output, - offsets_offset + container * sizeof(std::uint32_t), - static_cast(header_size + metadata[container].payload_offset)); - } + write_u32(output, + offsets_offset + container * sizeof(std::uint32_t), + static_cast(header_size + metadata[container].payload_offset)); } } @@ -750,34 +615,10 @@ __global__ void encode_sparse_row_kernel(std::uint32_t const* ids, auto const begin = metadata[container].begin; auto const end = container + 1 < num_containers ? metadata[container + 1].begin : cardinality; auto const payload = header_size + metadata[container].payload_offset; - if (metadata[container].kind == container_kind::array) { - for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { - write_u16(output, - payload + static_cast(i - begin) * sizeof(std::uint16_t), - static_cast(ids[i] & 0xffffu)); - } - continue; - } - - if (threadIdx.x == 0) { - write_u16(output, payload, metadata[container].runs); - std::uint16_t run_index{}; - for (auto i = begin; i < end;) { - auto const start = static_cast(ids[i] & 0xffffu); - auto j = i + 1; - while (j < end && - static_cast(ids[j]) == static_cast(ids[j - 1]) + 1) { - ++j; - } - auto const last = static_cast(ids[j - 1] & 0xffffu); - auto const run_offset = payload + sizeof(std::uint16_t) + - static_cast(run_index) * 2 * sizeof(std::uint16_t); - write_u16(output, run_offset, start); - write_u16( - output, run_offset + sizeof(std::uint16_t), static_cast(last - start)); - ++run_index; - i = j; - } + for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { + write_u16(output, + payload + static_cast(i - begin) * sizeof(std::uint16_t), + static_cast(ids[i] & 0xffffu)); } } __syncthreads(); @@ -841,8 +682,7 @@ device_build_result build_sparse_from_device_ids( RAFT_EXPECTS(host_summary.cardinality > 0 && host_summary.num_containers > 0, "Internal error: nonempty sparse Roaring input produced an empty device build."); - auto const header_size = - portable_header_size(host_summary.num_containers, host_summary.has_run != 0); + auto const header_size = portable_header_size(host_summary.num_containers); auto const serialized_bytes = header_size + static_cast(host_summary.payload_bytes); rmm::device_uvector output(0, stream); { @@ -859,22 +699,37 @@ device_build_result build_sparse_from_device_ids( metadata, host_summary.num_containers, header_size, - host_summary.has_run != 0, output.data(), reinterpret_cast(output.data() + reference_offset(serialized_bytes))); RAFT_CUDA_TRY(cudaPeekAtLastError()); return {std::move(output), serialized_bytes, size, true}; } +#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER +__global__ void validate_cuco_input_kernel(std::uint32_t const* ids, + std::size_t size, + std::uint64_t dataset_rows, + std::uint32_t* invalid) +{ + auto const i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < size && static_cast(ids[i]) >= dataset_rows) { atomicExch(invalid, 1u); } +} + +__global__ void store_cuco_ref_kernel(ref_type ref, ref_type* output) +{ + if (threadIdx.x == 0) { ::new (static_cast(output)) ref_type{ref}; } +} +#endif + /** * Build a standard portable Roaring row from device IDs. * - * Very sparse rows use the single-CTA builder above. Larger rows use one ID - * array for device-wide radix sorting. Both pre-sorted paths - * use the caller's strictly increasing IDs directly and allocate no normalization ID array. - * Boundary scan data, CUB workspace, and O(min(input IDs, 65536)) container metadata remain - * cardinality-scaled. In particular, the builder never allocates storage proportional to - * `dataset_rows` bits. + * When cuCollections provides its raw-index factories, the one-row path delegates construction to + * `roaring_bitmap::from_indices` or `from_sorted_unique_indices`, retains that owner without + * copying its serialized payload, and initializes the cuVS device reference once. The code below + * those factories remains the compatibility implementation for the currently pinned cuco revision. + * Multi-row inputs never enter this function; they use the segmented cuVS builder. Neither path + * allocates storage proportional to `dataset_rows` bits. * * Pre-sorted ordering and uniqueness are unchecked caller promises. */ @@ -888,6 +743,41 @@ device_build_result build_from_device_ids( auto const stream = raft::resource::get_cuda_stream(res); auto const size = static_cast(ids.extent(0)); if (size == 0) { return {rmm::device_uvector(0, stream), 0, 0, false}; } +#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER + { + // PR #839 owns the serialized bytes and already caches parsed metadata in + // its host-side ref. Keep that allocation alive and materialize only the + // lightweight ref in cuVS device storage; the payload is never recopied. + rmm::device_uvector invalid(1, stream); + RAFT_CUDA_TRY(cudaMemsetAsync(invalid.data(), 0, sizeof(std::uint32_t), stream)); + validate_cuco_input_kernel<<>>( + ids.data_handle(), size, dataset_rows, invalid.data()); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + std::uint32_t host_invalid{}; + RAFT_CUDA_TRY(cudaMemcpyAsync( + &host_invalid, invalid.data(), sizeof(host_invalid), cudaMemcpyDeviceToHost, stream)); + + 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); + // Both PR #839 factories perform their exact-size readback after all prior + // stream work, so the validation result is ready without another sync. + RAFT_EXPECTS(host_invalid == 0, "Roaring allowlist ID must be smaller than dataset_rows."); + + auto owner = std::make_unique(std::move(bitmap)); + auto const serialized_bytes = static_cast(owner->size_bytes()); + auto const cardinality = static_cast(owner->size()); + rmm::device_uvector output(sizeof(ref_type), stream); + store_cuco_ref_kernel<<<1, 1, 0, stream>>>(owner->ref(), + reinterpret_cast(output.data())); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + return {std::move(output), serialized_bytes, cardinality, true, std::move(owner)}; + } +#endif auto const sparse_cutoff = pre_sorted ? kSparseBuilderMaxPreSortedIds : kSparseBuilderMaxIds; if (size <= sparse_cutoff) { return build_sparse_from_device_ids(res, dataset_rows, ids, pre_sorted); @@ -949,7 +839,6 @@ device_build_result build_from_device_ids( reinterpret_cast(scratch.data() + layout.payload_sizes_offset); auto* payload_offsets = reinterpret_cast(scratch.data() + layout.payload_offsets_offset); - auto* has_run = reinterpret_cast(scratch.data() + layout.has_run_offset); auto* device_summary = reinterpret_cast(scratch.data() + layout.summary_offset); auto* workspace = scratch.data() + layout.workspace_offset; @@ -992,12 +881,8 @@ device_build_result build_from_device_ids( "roaring_allowlist::container_analysis"); RAFT_CUDA_TRY( cudaMemsetAsync(payload_sizes, 0, max_containers * sizeof(std::uint64_t), stream)); - RAFT_CUDA_TRY(cudaMemsetAsync(has_run, 0, sizeof(std::uint32_t), stream)); - analyze_containers_kernel<<(max_containers), - kBuilderBlockSize, - 0, - stream>>>( - normalized_ids, valid_count, container_starts, num_containers, kinds, payload_sizes, has_run); + analyze_containers_kernel<<>>( + valid_count, container_starts, num_containers, kinds, payload_sizes); RAFT_CUDA_TRY(cudaPeekAtLastError()); RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(workspace, payload_scan_workspace_bytes, @@ -1005,13 +890,8 @@ device_build_result build_from_device_ids( payload_offsets, container_slots, stream)); - finish_device_analysis_kernel<<<1, 1, 0, stream>>>(id_count, - valid_count, - num_containers, - has_run, - payload_sizes, - payload_offsets, - device_summary); + finish_device_analysis_kernel<<<1, 1, 0, stream>>>( + id_count, valid_count, num_containers, payload_sizes, payload_offsets, device_summary); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -1026,7 +906,7 @@ device_build_result build_from_device_ids( RAFT_EXPECTS(summary.cardinality > 0 && summary.num_containers > 0, "Internal error: nonempty Roaring input produced an empty device build."); - auto const header_size = portable_header_size(summary.num_containers, summary.has_run != 0); + auto const header_size = portable_header_size(summary.num_containers); RAFT_EXPECTS(summary.payload_bytes <= std::numeric_limits::max() - header_size, "Portable Roaring row exceeds the 32-bit offset range."); auto const serialized_bytes = header_size + static_cast(summary.payload_bytes); @@ -1036,8 +916,7 @@ device_build_result build_from_device_ids( "roaring_allowlist::final_allocation"); output.resize(owned_storage_bytes(serialized_bytes), stream); } - auto const header_items = - std::max(summary.num_containers, (summary.num_containers + 7) / 8); + auto const header_items = static_cast(summary.num_containers); { common::nvtx::range stage_scope("roaring_allowlist::header_encode"); encode_header_kernel<<>>( @@ -1045,10 +924,8 @@ device_build_result build_from_device_ids( valid_count, container_starts, summary.num_containers, - kinds, payload_offsets, header_size, - summary.has_run != 0, output.data()); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -1060,7 +937,6 @@ device_build_result build_from_device_ids( container_starts, summary.num_containers, kinds, - payload_sizes, payload_offsets, header_size, output.data()); @@ -1076,6 +952,9 @@ struct batched_device_build_result { std::vector serialized_bytes; std::vector cardinalities; bool references_initialized{}; +#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER + std::unique_ptr cuco_owner{}; +#endif }; struct packed_rows_layout { @@ -1133,7 +1012,6 @@ struct batch_general_scratch_layout { kinds_offset = reserve(max_containers, sizeof(container_kind), alignof(container_kind)); payload_sizes_offset = reserve(max_containers, sizeof(std::uint64_t), alignof(std::uint64_t)); payload_offsets_offset = reserve(max_containers, sizeof(std::uint64_t), alignof(std::uint64_t)); - has_run_offset = reserve(rows, sizeof(std::uint32_t), alignof(std::uint32_t)); summaries_offset = reserve(rows, sizeof(device_build_summary), alignof(device_build_summary)); output_offsets_offset = reserve(rows, sizeof(std::uint64_t), alignof(std::uint64_t)); workspace_offset = reserve(workspace_bytes, sizeof(cuda::std::byte), alignof(std::max_align_t)); @@ -1150,7 +1028,6 @@ struct batch_general_scratch_layout { std::size_t kinds_offset{}; std::size_t payload_sizes_offset{}; std::size_t payload_offsets_offset{}; - std::size_t has_run_offset{}; std::size_t summaries_offset{}; std::size_t output_offsets_offset{}; std::size_t workspace_offset{}; @@ -1233,50 +1110,32 @@ __global__ void fill_container_rows_kernel(std::int64_t const* row_container_off } } -__global__ void analyze_batch_containers_kernel(std::uint32_t const* ids, - std::int64_t const* indptr, +__global__ void analyze_batch_containers_kernel(std::int64_t const* indptr, std::int64_t const* valid_counts, std::int64_t const* container_starts, std::int64_t const* selected_count, std::uint32_t const* container_rows, container_kind* kinds, - std::uint64_t* payload_sizes, - std::uint32_t* has_run) + std::uint64_t* payload_sizes) { - auto const container = static_cast(blockIdx.x); - auto const count = *selected_count; - if (container >= count) { return; } - auto const row = static_cast(container_rows[container]); - auto const begin = container_starts[container]; - auto const row_end = indptr[row] + valid_counts[row]; - auto const end = container + 1 < count && container_rows[container + 1] == row - ? container_starts[container + 1] - : row_end; - std::uint32_t local_runs{}; - for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { - local_runs += - i == begin || static_cast(ids[i]) != static_cast(ids[i - 1]) + 1 - ? 1u - : 0u; - } - using block_reduce = cub::BlockReduce; - __shared__ typename block_reduce::TempStorage reduction_storage; - auto const runs = block_reduce(reduction_storage).Sum(local_runs); - if (threadIdx.x != 0) { return; } - auto const cardinality = static_cast(end - begin); - auto const normal_size = - cardinality <= kArrayCardinality ? cardinality * sizeof(std::uint16_t) : kBitmapBytes; - auto const run_size = sizeof(std::uint16_t) + runs * 2 * sizeof(std::uint16_t); - if (run_size < normal_size) { - kinds[container] = container_kind::run; - payload_sizes[container] = run_size; - atomicExch(has_run + row, 1u); - } else if (cardinality <= kArrayCardinality) { - kinds[container] = container_kind::array; - payload_sizes[container] = normal_size; - } else { - kinds[container] = container_kind::bitmap; - payload_sizes[container] = normal_size; + auto container = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + auto const stride = static_cast(gridDim.x) * blockDim.x; + auto const count = *selected_count; + for (; container < count; container += stride) { + auto const row = static_cast(container_rows[container]); + auto const begin = container_starts[container]; + auto const row_end = indptr[row] + valid_counts[row]; + auto const end = container + 1 < count && container_rows[container + 1] == row + ? container_starts[container + 1] + : row_end; + auto const cardinality = static_cast(end - begin); + if (cardinality <= kArrayCardinality) { + kinds[container] = container_kind::array; + payload_sizes[container] = cardinality * sizeof(std::uint16_t); + } else { + kinds[container] = container_kind::bitmap; + payload_sizes[container] = kBitmapBytes; + } } } @@ -1284,7 +1143,6 @@ __global__ void finish_batch_rows_kernel(std::int64_t rows, std::int64_t const* indptr, std::int64_t const* valid_counts, std::int64_t const* row_container_offsets, - std::uint32_t const* has_run, std::uint64_t const* payload_sizes, std::uint64_t const* payload_offsets, device_build_summary* summaries) @@ -1297,7 +1155,6 @@ __global__ void finish_batch_rows_kernel(std::int64_t rows, auto& summary = summaries[row]; summary.cardinality = cardinality; summary.num_containers = static_cast(end - begin); - summary.has_run = has_run[row]; summary.payload_bytes = begin == end ? 0 : payload_offsets[end - 1] + payload_sizes[end - 1] - payload_offsets[begin]; summary.invalid = valid_counts[row] != cardinality; @@ -1308,7 +1165,6 @@ __global__ void encode_batch_headers_kernel(std::uint32_t const* ids, std::int64_t const* valid_counts, std::int64_t const* container_starts, std::int64_t const* row_container_offsets, - container_kind const* kinds, std::uint64_t const* payload_offsets, device_build_summary const* summaries, std::uint64_t const* output_offsets, @@ -1317,34 +1173,15 @@ __global__ void encode_batch_headers_kernel(std::uint32_t const* ids, auto const row = static_cast(blockIdx.x); auto const summary = summaries[row]; if (summary.num_containers == 0) { return; } - auto const first_container = row_container_offsets[row]; - auto* output = storage + output_offsets[row]; - auto const has_run = summary.has_run != 0; - auto const run_bitmap_bytes = has_run ? (summary.num_containers + 7) / 8 : 0; - auto const descriptor_offset = - has_run ? sizeof(std::uint32_t) + run_bitmap_bytes : 2 * sizeof(std::uint32_t); + auto const first_container = row_container_offsets[row]; + auto* output = storage + output_offsets[row]; + constexpr auto descriptor_offset = 2 * sizeof(std::uint32_t); auto const offsets_offset = descriptor_offset + summary.num_containers * 2 * sizeof(std::uint16_t); - auto const header_size = portable_header_size(summary.num_containers, has_run); - bool const store_offsets = !has_run || summary.num_containers >= kOffsetThreshold; + auto const header_size = portable_header_size(summary.num_containers); if (threadIdx.x == 0) { - if (has_run) { - write_u32(output, 0, kCookieRun | ((summary.num_containers - 1) << 16)); - } else { - write_u32(output, 0, kCookieNoRun); - write_u32(output, sizeof(std::uint32_t), summary.num_containers); - } - } - auto* output_bytes = reinterpret_cast(output); - for (std::uint32_t byte = threadIdx.x; byte < run_bitmap_bytes; byte += blockDim.x) { - std::uint8_t value{}; - for (std::uint32_t bit = 0; bit < 8; ++bit) { - auto const local = byte * 8 + bit; - if (local < summary.num_containers && kinds[first_container + local] == container_kind::run) { - value |= static_cast(1u << bit); - } - } - output_bytes[sizeof(std::uint32_t) + byte] = value; + write_u32(output, 0, kCookieNoRun); + write_u32(output, sizeof(std::uint32_t), summary.num_containers); } auto const row_end = indptr[row] + valid_counts[row]; for (std::uint32_t local = threadIdx.x; local < summary.num_containers; local += blockDim.x) { @@ -1355,12 +1192,10 @@ __global__ void encode_batch_headers_kernel(std::uint32_t const* ids, write_u16(output, descriptor, static_cast(ids[begin] >> 16)); write_u16( output, descriptor + sizeof(std::uint16_t), static_cast(end - begin - 1)); - if (store_offsets) { - auto const local_payload = payload_offsets[container] - payload_offsets[first_container]; - write_u32(output, - offsets_offset + local * sizeof(std::uint32_t), - static_cast(header_size + local_payload)); - } + auto const local_payload = payload_offsets[container] - payload_offsets[first_container]; + write_u32(output, + offsets_offset + local * sizeof(std::uint32_t), + static_cast(header_size + local_payload)); } } @@ -1372,7 +1207,6 @@ __global__ void encode_batch_payloads_kernel(std::uint32_t const* ids, std::uint32_t const* container_rows, std::int64_t num_containers, container_kind const* kinds, - std::uint64_t const* payload_sizes, std::uint64_t const* payload_offsets, device_build_summary const* summaries, std::uint64_t const* output_offsets, @@ -1388,8 +1222,8 @@ __global__ void encode_batch_payloads_kernel(std::uint32_t const* ids, auto const row_end = indptr[row] + valid_counts[row]; auto const end = local + 1 < summary.num_containers ? container_starts[container + 1] : row_end; auto* output = storage + output_offsets[row]; - auto const payload = portable_header_size(summary.num_containers, summary.has_run != 0) + - payload_offsets[container] - payload_offsets[first_container]; + auto const payload = portable_header_size(summary.num_containers) + payload_offsets[container] - + payload_offsets[first_container]; if (kinds[container] == container_kind::array) { for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { @@ -1399,68 +1233,21 @@ __global__ void encode_batch_payloads_kernel(std::uint32_t const* ids, } return; } - using run_scan = cub::BlockScan; - union payload_scratch { - std::uint64_t bitmap_words[kBitmapBytes / sizeof(std::uint64_t)]; - typename run_scan::TempStorage run_scan_storage; - }; - __shared__ payload_scratch scratch; - __shared__ std::uint32_t run_base; - __shared__ std::uint32_t tile_runs; - if (kinds[container] == container_kind::bitmap) { - constexpr std::uint32_t words = kBitmapBytes / sizeof(std::uint64_t); - for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { - scratch.bitmap_words[word] = 0; - } - __syncthreads(); - for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { - auto const lower = ids[i] & 0xffffu; - atomicOr(reinterpret_cast(&scratch.bitmap_words[lower / 64]), - static_cast(std::uint64_t{1} << (lower % 64))); - } - __syncthreads(); - for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { - write_u64(output, payload + word * sizeof(std::uint64_t), scratch.bitmap_words[word]); - } - return; + + __shared__ std::uint64_t bitmap_words[kBitmapBytes / sizeof(std::uint64_t)]; + constexpr std::uint32_t words = kBitmapBytes / sizeof(std::uint64_t); + for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { + bitmap_words[word] = 0; } - auto const num_runs = static_cast( - (payload_sizes[container] - sizeof(std::uint16_t)) / (2 * sizeof(std::uint16_t))); - if (threadIdx.x == 0) { - write_u16(output, payload, num_runs); - run_base = 0; + __syncthreads(); + for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { + auto const lower = ids[i] & 0xffffu; + atomicOr(reinterpret_cast(&bitmap_words[lower / 64]), + static_cast(std::uint64_t{1} << (lower % 64))); } __syncthreads(); - for (auto tile = begin; tile < end; tile += blockDim.x) { - auto const i = tile + threadIdx.x; - std::uint32_t const is_run_start = - i < end && (i == begin || - static_cast(ids[i]) != static_cast(ids[i - 1]) + 1) - ? 1u - : 0u; - std::uint32_t run_rank{}; - std::uint32_t block_runs{}; - run_scan(scratch.run_scan_storage).ExclusiveSum(is_run_start, run_rank, block_runs); - if (threadIdx.x == 0) { tile_runs = block_runs; } - __syncthreads(); - if (is_run_start != 0) { - auto j = i + 1; - while (j < end && - static_cast(ids[j]) == static_cast(ids[j - 1]) + 1) { - ++j; - } - auto const start = static_cast(ids[i] & 0xffffu); - auto const last = static_cast(ids[j - 1] & 0xffffu); - auto const run_offset = - payload + sizeof(std::uint16_t) + - static_cast(run_base + run_rank) * 2 * sizeof(std::uint16_t); - write_u16(output, run_offset, start); - write_u16( - output, run_offset + sizeof(std::uint16_t), static_cast(last - start)); - } - __syncthreads(); - if (threadIdx.x == 0) { run_base += tile_runs; } - __syncthreads(); + for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { + write_u64(output, payload + word * sizeof(std::uint64_t), bitmap_words[word]); } } @@ -1558,7 +1345,6 @@ batched_device_build_result build_general_rows( reinterpret_cast(scratch.data() + layout.payload_sizes_offset); auto* payload_offsets = reinterpret_cast(scratch.data() + layout.payload_offsets_offset); - auto* has_run = reinterpret_cast(scratch.data() + layout.has_run_offset); auto* summaries = reinterpret_cast(scratch.data() + layout.summaries_offset); auto* output_offsets = @@ -1603,19 +1389,14 @@ batched_device_build_result build_general_rows( row_container_offsets, row_count, container_rows); RAFT_CUDA_TRY(cudaPeekAtLastError()); RAFT_CUDA_TRY(cudaMemsetAsync(payload_sizes, 0, max_containers * sizeof(std::uint64_t), stream)); - RAFT_CUDA_TRY(cudaMemsetAsync(has_run, 0, rows * sizeof(std::uint32_t), stream)); - analyze_batch_containers_kernel<<(max_containers), - kBuilderBlockSize, - 0, - stream>>>(normalized, - indptr.data_handle(), - valid_counts, - container_starts, - selected_count, - container_rows, - kinds, - payload_sizes, - has_run); + analyze_batch_containers_kernel<<>>( + indptr.data_handle(), + valid_counts, + container_starts, + selected_count, + container_rows, + kinds, + payload_sizes); RAFT_CUDA_TRY(cudaPeekAtLastError()); RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(workspace, scan_workspace_bytes, @@ -1628,7 +1409,6 @@ batched_device_build_result build_general_rows( indptr.data_handle(), valid_counts, row_container_offsets, - has_run, payload_sizes, payload_offsets, summaries); @@ -1653,7 +1433,7 @@ batched_device_build_result build_general_rows( "Internal error: general batched row analysis failed."); cardinalities[row] = cardinality; if (cardinality == 0) { continue; } - auto const header = portable_header_size(summary.num_containers, summary.has_run != 0); + auto const header = portable_header_size(summary.num_containers); RAFT_EXPECTS(summary.payload_bytes <= std::numeric_limits::max() - header, "Portable Roaring row exceeds the 32-bit offset range."); serialized[row] = header + static_cast(summary.payload_bytes); @@ -1673,7 +1453,6 @@ batched_device_build_result build_general_rows( valid_counts, container_starts, row_container_offsets, - kinds, payload_offsets, summaries, output_offsets, @@ -1690,7 +1469,6 @@ batched_device_build_result build_general_rows( container_rows, static_cast(actual_containers), kinds, - payload_sizes, payload_offsets, summaries, output_offsets, @@ -1748,12 +1526,22 @@ batched_device_build_result build_batched_from_device_ids( auto row_view = raft::make_device_vector_view( ids.data_handle(), static_cast(size)); auto built = build_from_device_ids(res, dataset_rows, row_view, pre_sorted); +#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER + return {std::move(built.storage), + {0}, + {0}, + {built.serialized_bytes}, + {built.cardinality}, + built.reference_initialized, + std::move(built.cuco_owner)}; +#else return {std::move(built.storage), {0}, {built.cardinality == 0 ? 0 : reference_offset(built.serialized_bytes)}, {built.serialized_bytes}, {built.cardinality}, built.reference_initialized}; +#endif } return build_general_rows(res, dataset_rows, ids, indptr, host_indptr, pre_sorted); } @@ -1972,6 +1760,9 @@ __global__ void contains_kernel(ref_type const* const* references, } // namespace struct roaring_allowlist::impl { +#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER + std::unique_ptr cuco_owner; +#endif rmm::device_uvector storage; std::vector row_offsets_; std::vector reference_offsets_; @@ -1984,7 +1775,12 @@ struct roaring_allowlist::impl { bool references_initialized_{}; impl(raft::resources const& res, std::size_t dataset_rows, batched_device_build_result&& built) +#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER + : cuco_owner(std::move(built.cuco_owner)), + storage(std::move(built.storage)), +#else : storage(std::move(built.storage)), +#endif row_offsets_(std::move(built.row_offsets)), reference_offsets_(std::move(built.reference_offsets)), serialized_bytes_(std::move(built.serialized_bytes)), @@ -2198,9 +1994,13 @@ std::size_t roaring_allowlist::total_cardinality() const noexcept std::size_t roaring_allowlist::size_bytes() const noexcept { - return impl_->storage.size() * sizeof(cuda::std::byte) + - impl_->references.size() * sizeof(ref_type const*) + - impl_->empty_rows.size() * sizeof(std::uint8_t); + auto bytes = impl_->storage.size() * sizeof(cuda::std::byte) + + impl_->references.size() * sizeof(ref_type const*) + + impl_->empty_rows.size() * sizeof(std::uint8_t); +#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER + if (impl_->cuco_owner) { bytes += static_cast(impl_->cuco_owner->size_bytes()); } +#endif + return bytes; } roaring_allowlist_view roaring_allowlist::view(std::size_t allowlist_id) const diff --git a/cpp/tests/neighbors/roaring_allowlist.cu b/cpp/tests/neighbors/roaring_allowlist.cu index b744af3935..df5662d98b 100644 --- a/cpp/tests/neighbors/roaring_allowlist.cu +++ b/cpp/tests/neighbors/roaring_allowlist.cu @@ -8,9 +8,6 @@ #include -#include "../../bench/ann/src/cuvs/allowlist_benchmark_utils.hpp" -#include "../../bench/ann/src/cuvs/cagra_filter_benchmark_utils.hpp" - #include #include #include @@ -40,9 +37,9 @@ namespace { // These hand-built fixtures exercise import of the portable bytes themselves. // Keep their field order aligned with // https://github.com/RoaringBitmap/RoaringFormatSpec. -constexpr std::uint32_t kCookieNoRun = 12346; -constexpr std::uint32_t kCookieRun = 12347; -constexpr std::uint32_t kSingleArrayPayloadByteOffset = 16; +constexpr std::uint32_t kCookieNoRun = 12346; +constexpr std::uint32_t kCookieRun = 12347; +constexpr std::uint32_t kSingleContainerPayloadByteOffset = 16; void append_u16(std::vector& out, std::uint16_t value) { @@ -57,6 +54,16 @@ void append_u32(std::vector& out, std::uint32_t value) } } +std::uint32_t read_u32(std::vector const& bytes, std::size_t offset = 0) +{ + std::uint32_t value{}; + for (int i = 0; i < 4; ++i) { + value |= static_cast(std::to_integer(bytes[offset + i])) + << (8 * i); + } + return value; +} + // One no-run array container: cookie, N, {key, cardinality - 1}, payload // offset, values. std::vector array_row(std::vector const& values) @@ -66,7 +73,7 @@ std::vector array_row(std::vector const& values) append_u32(out, 1); append_u16(out, 0); append_u16(out, static_cast(values.size() - 1)); - append_u32(out, kSingleArrayPayloadByteOffset); + append_u32(out, kSingleContainerPayloadByteOffset); for (auto value : values) { append_u16(out, value); } @@ -89,6 +96,23 @@ std::vector run_row(std::uint16_t start, std::uint16_t length_minus_o return out; } +// One no-run bitmap container containing the 5,000 even values below 10,000. +std::vector bitmap_row() +{ + constexpr std::uint16_t cardinality = 5000; + std::vector out; + append_u32(out, kCookieNoRun); + append_u32(out, 1); + append_u16(out, 0); + append_u16(out, cardinality - 1); + append_u32(out, kSingleContainerPayloadByteOffset); + out.resize(kSingleContainerPayloadByteOffset + 8192, std::byte{0}); + for (std::uint32_t value = 0; value < 10000; value += 2) { + out[kSingleContainerPayloadByteOffset + value / 8] |= static_cast(1u << (value % 8)); + } + return out; +} + roaring_allowlist from_ragged_ids(raft::resources const& res, std::size_t dataset_rows, std::vector const& ids, @@ -241,7 +265,7 @@ void expect_batch_membership(raft::resources const& res, EXPECT_EQ(actual, expected); } -TEST(RoaringAllowlist, BuildsArrayRunBitmapAndMultiContainerRowsFromIds) +TEST(RoaringAllowlist, BuildsArrayBitmapAndMultiContainerRowsFromIds) { raft::device_resources res; @@ -256,9 +280,9 @@ TEST(RoaringAllowlist, BuildsArrayRunBitmapAndMultiContainerRowsFromIds) for (std::uint32_t id = 100; id < 300; ++id) { consecutive.push_back(id); } - auto run = from_ids(res, 1000, consecutive); - EXPECT_EQ(run.cardinality(0), 200); - expect_membership(res, run, {99, 100, 199, 299, 300}, {0, 1, 1, 1, 0}); + auto contiguous = from_ids(res, 1000, consecutive); + EXPECT_EQ(contiguous.cardinality(0), 200); + expect_membership(res, contiguous, {99, 100, 199, 299, 300}, {0, 1, 1, 1, 0}); std::vector sparse; for (std::uint32_t id = 0; id < 10000; id += 2) { @@ -307,6 +331,7 @@ TEST(RoaringAllowlist, BuildsRaggedRowsInOneGeneralBatch) EXPECT_NE(allowlists.view(0).device_reference(), allowlists.view(2).device_reference()); EXPECT_EQ(device_allowlists.total_cardinality(), ids.size()); EXPECT_EQ(device_allowlists.size_bytes(), allowlists.size_bytes()); + EXPECT_EQ(read_u32(copy_serialized_bytes(res, allowlists)), kCookieNoRun); expect_batch_membership(res, allowlists, @@ -421,7 +446,7 @@ TEST(RoaringAllowlist, ImportsRaggedPortableRowsIntoOnePackedOwner) {1, 0, 1, 1, 0, 0, 0, 0, 0, 1, 1, 0}); } -TEST(RoaringAllowlist, ConstructionPathsEmitByteIdenticalPortableRows) +TEST(RoaringAllowlist, CucoConstructionPathsEmitByteIdenticalPortableRows) { raft::device_resources res; constexpr std::uint32_t second_key = std::uint32_t{1} << 16; @@ -439,53 +464,17 @@ TEST(RoaringAllowlist, ConstructionPathsEmitByteIdenticalPortableRows) auto sorted = ids; std::sort(sorted.begin(), sorted.end()); auto const dataset_rows = static_cast(third_key) + 10000; - auto const expected = cuvs::bench::allowlist::make_portable_roaring_bytes(ids); - - auto host_unsorted = from_ids(res, dataset_rows, ids); - auto host_pre_sorted = from_ids(res, dataset_rows, sorted, true); - auto device_unsorted = from_device_ids(res, dataset_rows, ids); - auto device_sorted = from_device_ids(res, dataset_rows, sorted, true); + auto host_unsorted = from_ids(res, dataset_rows, ids); + auto host_pre_sorted = from_ids(res, dataset_rows, sorted, true); + auto device_unsorted = from_device_ids(res, dataset_rows, ids); + auto device_sorted = from_device_ids(res, dataset_rows, sorted, true); - EXPECT_EQ(copy_serialized_bytes(res, host_unsorted), expected); + auto const expected = copy_serialized_bytes(res, host_unsorted); EXPECT_EQ(copy_serialized_bytes(res, host_pre_sorted), expected); EXPECT_EQ(copy_serialized_bytes(res, device_unsorted), expected); EXPECT_EQ(copy_serialized_bytes(res, device_sorted), expected); } -TEST(RoaringAllowlist, BenchmarkPortableSerializerImportsMixedContainerKinds) -{ - raft::device_resources res; - constexpr std::uint32_t second_key = std::uint32_t{1} << 16; - constexpr std::uint32_t third_key = std::uint32_t{2} << 16; - - std::vector ids; - for (std::uint32_t id = 100; id < 300; ++id) { - ids.push_back(id); // run - } - ids.insert(ids.end(), {second_key + 1, second_key + 3, second_key + 5}); // array - for (std::uint32_t id = 0; id < 10000; id += 2) { - ids.push_back(third_key + id); // bitmap - } - std::reverse(ids.begin(), ids.end()); - auto const bytes = cuvs::bench::allowlist::make_portable_roaring_bytes(ids); - auto imported = import_row(res, static_cast(third_key) + 10000, bytes); - - EXPECT_EQ(imported.cardinality(0), ids.size()); - expect_membership(res, - imported, - {99, - 100, - 299, - 300, - second_key + 1, - second_key + 2, - third_key, - third_key + 1, - third_key + 9998, - third_key + 9999}, - {0, 1, 1, 0, 1, 0, 1, 0, 1, 0}); -} - TEST(RoaringAllowlist, BuildsFromUnsortedUniqueDeviceIdsWithoutModifyingInput) { raft::device_resources res; @@ -494,7 +483,7 @@ TEST(RoaringAllowlist, BuildsFromUnsortedUniqueDeviceIdsWithoutModifyingInput) std::vector ids; for (std::uint32_t id = 100; id < 300; ++id) { - ids.push_back(id); // run container + ids.push_back(id); // contiguous array container } ids.insert(ids.end(), {second_key + 1, second_key + 3, second_key + 5}); // array container for (std::uint32_t id = 0; id < 10000; id += 2) { @@ -581,7 +570,7 @@ TEST(RoaringAllowlist, SparseDeviceBuilderHandlesThresholdAndGeneralCrossover) std::vector sorted; sorted.reserve(129); for (std::uint32_t id = 100; id < 164; ++id) { - sorted.push_back(id); // one run container + sorted.push_back(id); // one contiguous array container } for (std::uint32_t value = 1; value < 65; value += 2) { sorted.push_back(second_key + value); // one array container @@ -709,6 +698,11 @@ TEST(RoaringAllowlist, ImportsPortableArrayRunAndEmptyRows) EXPECT_EQ(run.cardinality(0), 100); expect_membership(res, run, {99, 100, 199, 200}, {0, 1, 1, 0}, true); + auto bitmap_bytes = bitmap_row(); + auto bitmap = import_row(res, 10000, bitmap_bytes); + EXPECT_EQ(bitmap.cardinality(0), 5000); + expect_membership(res, bitmap, {0, 1, 8192, 9998, 9999}, {1, 0, 1, 1, 0}, true); + std::vector standard_empty; append_u32(standard_empty, kCookieNoRun); append_u32(standard_empty, 0); @@ -831,114 +825,5 @@ TEST(RoaringFilter, RejectsSmallInvalidMappings) EXPECT_THROW(filter.set_allowlist(res, 0, eleven.view(0)), raft::logic_error); } -TEST(CagraFilterBenchmarkUtils, GeneratesExactDeterministicAllowlistLayouts) -{ - using cuvs::bench::cagra_filter::allowlist_layout; - using cuvs::bench::cagra_filter::make_allowed_ids; - - for (auto layout : {allowlist_layout::random, allowlist_layout::runs}) { - auto first = make_allowed_ids(257, 73, layout, 12345); - auto second = make_allowed_ids(257, 73, layout, 12345); - EXPECT_EQ(first, second); - ASSERT_EQ(first.size(), 73); - - std::sort(first.begin(), first.end()); - EXPECT_EQ(std::adjacent_find(first.begin(), first.end()), first.end()); - EXPECT_LT(first.back(), 257); - } - - auto run = make_allowed_ids(257, 73, allowlist_layout::runs, 12345); - std::sort(run.begin(), run.end()); - auto breaks = std::size_t{0}; - for (std::size_t i = 1; i < run.size(); ++i) { - breaks += run[i] != run[i - 1] + 1 ? 1 : 0; - } - EXPECT_LE(breaks, - 1); // a circular interval is one range, or two when it wraps at row zero - - auto random = make_allowed_ids(257, 73, allowlist_layout::random, 12345); - std::sort(random.begin(), random.end()); - breaks = 0; - for (std::size_t i = 1; i < random.size(); ++i) { - breaks += random[i] != random[i - 1] + 1 ? 1 : 0; - } - EXPECT_GT(breaks, 1); - - EXPECT_TRUE(make_allowed_ids(257, 0, allowlist_layout::random, 12345).empty()); - EXPECT_THROW(make_allowed_ids(0, 0, allowlist_layout::random, 12345), std::invalid_argument); - EXPECT_THROW(make_allowed_ids(10, 11, allowlist_layout::runs, 12345), std::invalid_argument); -} - -TEST(CagraFilterBenchmarkUtils, SummarizesLatencySamples) -{ - auto const stats = cuvs::bench::cagra_filter::summarize({1.0, 2.0, 3.0, 4.0, 5.0}); - EXPECT_DOUBLE_EQ(stats.mean, 3.0); - EXPECT_DOUBLE_EQ(stats.p50, 3.0); - EXPECT_DOUBLE_EQ(stats.p95, 4.8); - EXPECT_DOUBLE_EQ(stats.p99, 4.96); - - EXPECT_EQ(cuvs::bench::cagra_filter::cardinality_for_fraction(1'000'000, 0.001, 10), 1000); - EXPECT_EQ(cuvs::bench::cagra_filter::allowed_percent_label(0.001), "allowed_0.1pct"); - EXPECT_EQ(cuvs::bench::cagra_filter::allowed_percent_label(0.015), "allowed_1.5pct"); -} - -TEST(AllowlistBenchmarkUtils, GeneratesControlledSmallLookupProfiles) -{ - using namespace cuvs::bench::allowlist; - - constexpr std::uint32_t rows = 64; - constexpr std::size_t cardinality = 16; - for (auto profile : {lookup_profile::best, lookup_profile::average, lookup_profile::worst}) { - auto first = make_allowed_ids(rows, cardinality, profile, 12345); - auto second = make_allowed_ids(rows, cardinality, profile, 12345); - EXPECT_EQ(first, second); - ASSERT_EQ(first.size(), cardinality); - std::sort(first.begin(), first.end()); - EXPECT_EQ(std::adjacent_find(first.begin(), first.end()), first.end()); - EXPECT_LT(first.back(), rows); - } - - auto best = make_allowed_ids(rows, cardinality, lookup_profile::best, 12345); - EXPECT_EQ(occupied_chunks(best), 1); - EXPECT_EQ(count_runs(best), 1); - - auto worst = make_allowed_ids(rows, cardinality, lookup_profile::worst, 12345); - EXPECT_EQ(maximum_encoded_runs(cardinality, rows), 7); - EXPECT_EQ(count_runs(worst), 7); - - EXPECT_THROW(make_allowed_ids(0, 0, lookup_profile::best, 12345), std::invalid_argument); - EXPECT_THROW(make_allowed_ids(10, 11, lookup_profile::worst, 12345), std::invalid_argument); -} - -TEST(AllowlistBenchmarkUtils, GeneratesBestHitsAverageDomainAndWorstMisses) -{ - using namespace cuvs::bench::allowlist; - constexpr std::uint32_t rows = 64; - constexpr std::size_t cardinality = 16; - for (auto profile : {lookup_profile::best, lookup_profile::average, lookup_profile::worst}) { - auto ids = make_allowed_ids(rows, cardinality, profile, 12345); - auto probes = make_probes(rows, ids, profile, 257, 67890); - ASSERT_EQ(probes.size(), 257); - EXPECT_TRUE(std::all_of(probes.begin(), probes.end(), [](auto probe) { return probe < rows; })); - std::sort(ids.begin(), ids.end()); - if (profile == lookup_profile::best) { - EXPECT_TRUE(std::all_of(probes.begin(), probes.end(), [&](auto probe) { - return std::binary_search(ids.begin(), ids.end(), probe); - })); - } else if (profile == lookup_profile::worst) { - EXPECT_TRUE(std::none_of(probes.begin(), probes.end(), [&](auto probe) { - return std::binary_search(ids.begin(), ids.end(), probe); - })); - } - } - EXPECT_THROW(make_probes(rows, {1}, lookup_profile::average, 0, 67890), std::invalid_argument); - - auto full_domain = make_allowed_ids(std::uint64_t{1} << 32, 1024, lookup_profile::average, 12345); - auto repeated = make_allowed_ids(std::uint64_t{1} << 32, 1024, lookup_profile::average, 12345); - EXPECT_EQ(full_domain, repeated); - std::sort(full_domain.begin(), full_domain.end()); - EXPECT_EQ(std::adjacent_find(full_domain.begin(), full_domain.end()), full_domain.end()); -} - } // namespace } // namespace cuvs::core From 5668d01e2d7688f1d8a7cf3c05769474813d7805 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Wed, 2 Sep 2026 20:29:38 +0000 Subject: [PATCH 3/5] just single batch for now --- cpp/CMakeLists.txt | 1 + cpp/cmake/patches/cuco_override.json | 9 + cpp/include/cuvs/core/roaring_allowlist.hpp | 109 +- cpp/include/cuvs/neighbors/common.hpp | 36 +- cpp/src/core/roaring_allowlist.cu | 1998 +---------------- .../neighbors/ann_cagra/test_filter_udf.cu | 108 +- cpp/tests/neighbors/roaring_allowlist.cu | 748 +----- .../cpp/src/cagra_roaring_filter_example.cu | 111 +- 8 files changed, 273 insertions(+), 2847 deletions(-) create mode 100644 cpp/cmake/patches/cuco_override.json diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 4193509d2a..d6cc530457 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -189,6 +189,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) diff --git a/cpp/cmake/patches/cuco_override.json b/cpp/cmake/patches/cuco_override.json new file mode 100644 index 0000000000..55204afb14 --- /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": "9d7c9307395c3b8795d93ad65d0751c98471dde6" + } + } +} diff --git a/cpp/include/cuvs/core/roaring_allowlist.hpp b/cpp/include/cuvs/core/roaring_allowlist.hpp index 60789d5956..bcf706989d 100644 --- a/cpp/include/cuvs/core/roaring_allowlist.hpp +++ b/cpp/include/cuvs/core/roaring_allowlist.hpp @@ -19,7 +19,7 @@ namespace CUVS_EXPORT cuvs { namespace core { /** - * @brief Non-owning device view of one row in a batched Roaring allowlist. + * @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 @@ -35,7 +35,7 @@ class CUVS_EXPORT roaring_allowlist_view { [[nodiscard]] bool empty() const noexcept { return cardinality_ == 0; } [[nodiscard]] bool valid() const noexcept { return valid_; } - /** @brief Opaque device pointer to the pre-parsed cuCollections reference, or null if empty. */ + /** @brief Opaque device pointer to the initialized cuCollections reference, or null if empty. */ [[nodiscard]] void const* device_reference() const noexcept { return device_reference_; } private: @@ -58,30 +58,23 @@ class CUVS_EXPORT roaring_allowlist_view { }; /** - * @brief Owning immutable batch of exact per-query Roaring allowlists. + * @brief Owning immutable exact Roaring allowlist over CAGRA dataset-row IDs. * - * Logically, the owner is a sparse matrix with one allowlist row per query and one possible column - * per dataset row. @ref from_ids accepts one contiguous ID vector plus an indptr vector that - * delimits independently sized query rows. Every row is sorted and encoded independently. For - * multiple rows, all variable-length portable Roaring streams and their initialized - * `cuco::experimental::roaring_bitmap_ref` objects share one packed device allocation. - * The batch builder uses indptr directly for segmented device radix sort and schedules - * analysis/encoding over all containers in all rows. + * 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. * - * A one-row input delegates raw-index construction to cuCollections. cuVS retains the cuco owner - * directly and materializes only its lightweight reference in cuVS device metadata; the serialized - * payload is not copied. cuco currently emits array and bitmap containers on this path. The cuVS - * multi-row builder emits the same array and bitmap container forms. Imported portable rows may - * still contain run containers. Final encoding and reference initialization remain stream ordered - * in both paths. + * 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. * - * IDs must be unique within each row. Setting @p pre_sorted skips sorting and promises that every - * row is strictly increasing; ordering and uniqueness are not checked. Every ID must be smaller - * than dataset_rows. + * 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/blob/9d7c9307395c3b8795d93ad65d0751c98471dde6/include/cuco/roaring_bitmap_ref.cuh * @see https://github.com/NVIDIA/cuCollections/pull/839 */ class CUVS_EXPORT roaring_allowlist { @@ -89,51 +82,29 @@ class CUVS_EXPORT roaring_allowlist { struct impl; public: - using key_type = std::uint32_t; - using indptr_type = std::int64_t; + using key_type = std::uint32_t; /** - * @brief Build ragged allowlist rows from contiguous host IDs and row offsets. + * @brief Build one allowlist from host IDs. * - * `indptr` contains `num_allowlists + 1` entries, starts at zero, is nondecreasing, - * and ends at `ids.extent(0)`. Empty slices are valid allowlists. + * Host IDs are copied to the construction stream and then use the same device builder as the + * device overload. Empty input is valid and rejects every candidate. */ static roaring_allowlist from_ids(raft::resources const& res, std::size_t dataset_rows, raft::host_vector_view ids, - raft::host_vector_view indptr, bool pre_sorted = false); /** - * @brief Build ragged allowlist rows from contiguous device IDs and row offsets. + * @brief Build one allowlist from device IDs. * - * The same indptr invariants as the host overload apply. The row offsets are copied to the host - * once for validation, shape-aware dispatch, and exact packed allocation. - * - * The input must remain valid until the construction stream reaches the work enqueued by this - * call. Temporary memory is O(total input IDs + total containers); no dense dataset-sized bitmap - * is materialized. - */ - static roaring_allowlist from_ids( - raft::resources const& res, - std::size_t dataset_rows, - raft::device_vector_view ids, - raft::device_vector_view indptr, - bool pre_sorted = false); - - /** - * @brief Import packed standard 32-bit portable Roaring rows. - * - * `byte_offsets` has `num_allowlists + 1` entries, starts at zero, is nondecreasing, and ends at - * `bytes.extent(0)`. Empty slices represent empty allowlists. Every row is strictly validated on - * the host before its bytes are copied. The host buffers must remain valid until the construction - * stream completes; pinned bytes are recommended when overlap matters. + * The input must remain valid until the construction stream reaches the enqueued work. + * Temporary memory is O(cardinality + container count); no dataset-sized dense bitmap is used. */ - static roaring_allowlist from_serialized( - raft::resources const& res, - std::size_t dataset_rows, - raft::host_vector_view bytes, - raft::host_vector_view byte_offsets); + 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(); @@ -142,33 +113,25 @@ class CUVS_EXPORT roaring_allowlist { roaring_allowlist(roaring_allowlist&&) noexcept; roaring_allowlist& operator=(roaring_allowlist&&) noexcept; - [[nodiscard]] std::size_t num_allowlists() const noexcept; [[nodiscard]] std::size_t dataset_rows() const noexcept; - [[nodiscard]] std::size_t cardinality(std::size_t allowlist_id) const; - [[nodiscard]] bool empty(std::size_t allowlist_id) const; - [[nodiscard]] std::size_t total_cardinality() const noexcept; + [[nodiscard]] std::size_t cardinality() const noexcept; + [[nodiscard]] bool empty() const noexcept; - /** @brief Total device bytes retained by packed rows, references, and row pointer metadata. */ + /** @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 of one row. */ - [[nodiscard]] roaring_allowlist_view view(std::size_t allowlist_id) const; + /** @brief Return a zero-copy view. */ + [[nodiscard]] roaring_allowlist_view view() const noexcept; - /** - * @brief Test a matrix of row IDs and synchronize the resource stream. - * - * `row_ids[q][i]` is tested against allowlist row `q`. Input and output shapes must match, and - * their first extent must equal @ref num_allowlists. - */ + /** @brief Test row IDs and synchronize the resource stream. */ void contains(raft::resources const& res, - raft::device_matrix_view row_ids, - raft::device_matrix_view output) const; + raft::device_vector_view row_ids, + raft::device_vector_view output) const; /** @brief Stream-ordered asynchronous version of @ref contains. */ - void contains_async( - raft::resources const& res, - raft::device_matrix_view row_ids, - raft::device_matrix_view output) const; + void contains_async(raft::resources const& res, + raft::device_vector_view row_ids, + raft::device_vector_view output) const; private: explicit roaring_allowlist(std::unique_ptr impl) noexcept; diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 4624889c5e..83ec04e345 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -1505,33 +1505,29 @@ struct bloom_filter : public base_filter { }; /** - * @brief Reusable per-query mapping to an immutable batch of exact Roaring allowlists. + * @brief Reusable per-query mapping to immutable exact Roaring allowlists. * - * Entry @c q selects row @c q of the owner. CAGRA retains candidate dataset row @c r when that - * allowlist contains @c r. Constructing from a @c cuvs::core::roaring_allowlist copies only its - * 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, parsing, - * initialization, synchronization, or per-query preprocessing. + * 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} - * // Flat IDs plus num_queries + 1 row offsets. - * auto allowlists = cuvs::core::roaring_allowlist::from_ids( + * auto first = cuvs::core::roaring_allowlist::from_ids( * res, dataset_rows, - * raft::make_host_vector_view(allowed_ids.data(), - * allowed_ids.size()), - * raft::make_host_vector_view(indptr.data(), - * indptr.size())); - * std::vector views; - * for (std::size_t q = 0; q < allowlists.num_allowlists(); ++q) { - * views.push_back(allowlists.view(q)); - * } + * 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 * - * The span overload remains useful when queries reuse rows from several owners or when one query's - * mapping must be replaced without rebuilding encoded allowlists. This filter owns its mapping - * tables and device payload, but not the referenced owner(s), which must outlive the filter and all - * searches using it. Copies are cheap shared handles required by CAGRA query-offset wrappers. + * 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. * * @see cuvs::core::roaring_allowlist * @see https://github.com/RoaringBitmap/RoaringFormatSpec diff --git a/cpp/src/core/roaring_allowlist.cu b/cpp/src/core/roaring_allowlist.cu index 141f561e87..33a3d5a6f9 100644 --- a/cpp/src/core/roaring_allowlist.cu +++ b/cpp/src/core/roaring_allowlist.cu @@ -1,30 +1,13 @@ /* * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * SPDX-License-Identifier: Apache-2.0 */ #include #include "nvtx.hpp" -#include - -// cuCollections PR #839 adds GPU construction from raw indices. Keep the -// existing cuVS builder as a compatibility fallback until that API reaches the -// pinned cuco revision, and keep using it for the segmented multi-row path. -#if __has_include() #include -#include -#define CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER 1 -#else -#define CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER 0 -#endif - -#include -#include -#include -#include -#include #include #include @@ -32,140 +15,51 @@ #include #include - -#include -#include +#include #include -#include +#include #include -#include #include #include #include #include -#include #include #include -#include namespace cuvs::core { namespace { /** - * Portable Roaring serialization used by each allowlist - * ===================================================== - * - * This file writes and validates the standard 32-bit portable Roaring format. - * The authoritative format description is: - * - * https://github.com/RoaringBitmap/RoaringFormatSpec - * - * The resulting bytes are consumed on the device by cuCollections: - * - * https://github.com/NVIDIA/cuCollections/blob/6001618aaa7f17ea2bbcd444650e9573c4f3d6c5/include/cuco/roaring_bitmap_ref.cuh - * https://github.com/NVIDIA/cuCollections/blob/6001618aaa7f17ea2bbcd444650e9573c4f3d6c5/include/cuco/detail/roaring_bitmap/util.cuh - * - * All integers below are little-endian. A 32-bit ID is divided into a container - * key and a value: - * - * ID = (uint32_t(key) << 16) | value - * high 16 bits low 16 bits - * - * IDs with the same key belong to one container. Container keys and array - * values are strictly increasing. The portable stream for one nonempty - * allowlist has one of these two headers: + * cuCollections Roaring ownership and encoding + * ============================================= * - * With no run containers, the row is laid out as: + * 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. * - * @code{.unparsed} - * uint32 cookie = 12346 - * uint32 N - * descriptor[N] - * uint32 container_offset[N] - * container payloads - * @endcode + * 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 materializes that reference in + * a stable device allocation during construction. Views and CAGRA filters copy only its pointer, so + * search never copies or reparses the payload. * - * Imported rows with at least one run container are laid out as: - * - * @code{.unparsed} - * uint32 cookie = 12347 | ((N - 1) << 16) - * uint8 run_container_bitmap[ceil(N / 8)] - * descriptor[N] - * uint32 container_offset[N] // present only when N >= 4 - * container payloads - * @endcode - * - * Each four-byte descriptor is: - * - * uint16 key - * uint16 cardinality_minus_one - * - * An offset is measured from the first byte of this stream. Without run - * containers, cardinality selects the payload representation: at most 4096 - * values use an array; more than 4096 use a bitmap. The run-container bitmap - * overrides that choice for marked containers. Its bit order is - * least-significant bit first. - * - * @code{.unparsed} - * array: - * uint16 value[cardinality] - * - * bitmap: - * uint64 words[1024] // 8192 bytes; v is bit (v % 64) of word (v / 64) - * - * run: - * uint16 number_of_runs - * { uint16 start; uint16 length_minus_one; } runs[number_of_runs] - * @endcode - * - * ID-based construction emits only the array and bitmap forms above. The run - * form is accepted only when importing an existing standard-portable stream. - * - * The portable format above describes exactly one bitmap. Query-to-allowlist - * association is a separate concern: `filtering::roaring_filter` stores device - * pointers to already initialized allowlist references. Consequently neither - * the serialized payload nor this metadata is copied or parsed by CAGRA search. + * @see https://github.com/RoaringBitmap/RoaringFormatSpec + * @see https://github.com/NVIDIA/cuCollections/pull/839 */ -// Names and thresholds used by the portable format specification. -constexpr std::uint32_t kCookieNoRun = 12346; -constexpr std::uint32_t kCookieRun = 12347; -constexpr std::size_t kArrayCardinality = 4096; -constexpr std::size_t kBitmapBytes = 8192; -constexpr std::size_t kOffsetThreshold = 4; - -using ref_type = cuco::experimental::roaring_bitmap_ref; - -#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER +using ref_type = cuco::experimental::roaring_bitmap_ref; using cuco_bitmap_allocator = rmm::mr::polymorphic_allocator; using cuco_bitmap_type = cuco::experimental::roaring_bitmap; -#endif - -struct row_metadata { - std::size_t cardinality{}; - std::uint32_t max_id{}; - bool empty{true}; -}; -std::uint16_t read_u16(std::byte const* data, std::size_t size, std::size_t offset) -{ - RAFT_EXPECTS(offset <= size && size - offset >= 2, - "Malformed portable Roaring bitmap: truncated uint16 value."); - return static_cast(std::to_integer(data[offset])) | - static_cast(std::to_integer(data[offset + 1])) << 8; -} +constexpr int kBlockSize = 256; -std::uint32_t read_u32(std::byte const* data, std::size_t size, std::size_t offset) +int grid_size_for(std::size_t count) { - RAFT_EXPECTS(offset <= size && size - offset >= 4, - "Malformed portable Roaring bitmap: truncated uint32 value."); - std::uint32_t value{}; - for (int i = 0; i < 4; ++i) { - value |= static_cast(std::to_integer(data[offset + i])) << (8 * i); - } - return value; + auto const blocks = (count + kBlockSize - 1) / kBlockSize; + return static_cast(std::min(blocks, 65535)); } void validate_dataset_rows(std::size_t dataset_rows) @@ -176,563 +70,39 @@ void validate_dataset_rows(std::size_t dataset_rows) "dataset_rows exceeds the uint32_t Roaring key domain."); } -enum class container_kind : std::uint8_t { array, bitmap }; - -std::size_t align_up(std::size_t offset, std::size_t alignment) -{ - return (offset + alignment - 1) / alignment * alignment; -} - -struct device_build_summary { - std::int64_t cardinality{}; - std::uint64_t payload_bytes{}; - std::uint32_t num_containers{}; - std::uint32_t invalid{}; -}; - -struct device_build_result { - rmm::device_uvector storage; - std::size_t serialized_bytes{}; - std::size_t cardinality{}; - bool reference_initialized{}; -#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER - std::unique_ptr cuco_owner{}; -#endif -}; - -std::size_t reference_offset(std::size_t serialized_bytes) -{ - return align_up(serialized_bytes, alignof(ref_type)); -} - -std::size_t owned_storage_bytes(std::size_t serialized_bytes) +__global__ void validate_input_kernel(std::uint32_t const* ids, + std::size_t size, + std::uint64_t dataset_rows, + std::uint32_t* invalid) { - return serialized_bytes == 0 ? 0 : reference_offset(serialized_bytes) + sizeof(ref_type); -} - -constexpr int kBuilderBlockSize = 256; - -// Small allowlists do not benefit from the general builder's device-wide sort, -// two scans, and separate per-stage allocations. At this cardinality every -// portable container is necessarily an array (a bitmap requires more than 4096 values in one -// high-16-bit partition), so one CTA can sort, analyze, and later encode the complete row. For a -// single pre-sorted row the cutoff is lower because the general path already avoids its most -// expensive stage, the device-wide sort. Batched rows use the 128-ID capacity because the launch is -// amortized across the matrix; keep both rules tied to the construction benchmark. -constexpr int kSparseBuilderBlockSize = 128; -constexpr std::size_t kSparseBuilderMaxIds = 128; -constexpr std::size_t kSparseBuilderMaxPreSortedIds = 64; -constexpr int kSparseItemsPerThread = - static_cast(kSparseBuilderMaxIds) / kSparseBuilderBlockSize; -static_assert(kSparseBuilderMaxIds % kSparseBuilderBlockSize == 0); - -struct sparse_container_metadata { - std::uint32_t begin{}; - std::uint32_t payload_offset{}; -}; - -static_assert(sizeof(sparse_container_metadata) == 8); - -struct sparse_scratch_layout { - explicit sparse_scratch_layout(std::size_t ids, std::size_t containers, bool store_sorted_ids) - { - sorted_ids_offset = align_up(sizeof(device_build_summary), alignof(std::uint32_t)); - auto offset = sorted_ids_offset + (store_sorted_ids ? ids * sizeof(std::uint32_t) : 0); - metadata_offset = align_up(offset, alignof(sparse_container_metadata)); - bytes = metadata_offset + containers * sizeof(sparse_container_metadata); - } - - std::size_t sorted_ids_offset{}; - std::size_t metadata_offset{}; - std::size_t bytes{}; -}; - -/** - * One allocation for all general-builder temporaries. - * - * The allocation is cardinality/container scaled. The largest CUB workspace is reused by sort, - * boundary selection, and payload scan because those stages are stream ordered. - */ -struct general_scratch_layout { - general_scratch_layout(std::size_t ids, - std::size_t containers, - bool store_sorted_ids, - std::size_t workspace_bytes) - { - std::size_t cursor{}; - auto reserve = [&](std::size_t count, std::size_t item_size, std::size_t alignment) { - auto const result = align_up(cursor, alignment); - cursor = result + count * item_size; - return result; - }; - - if (store_sorted_ids) { - sorted_ids_offset = reserve(ids, sizeof(std::uint32_t), alignof(std::uint32_t)); - } - id_count_offset = reserve(1, sizeof(std::int64_t), alignof(std::int64_t)); - valid_count_offset = reserve(1, sizeof(std::int64_t), alignof(std::int64_t)); - selected_count_offset = reserve(1, sizeof(std::int64_t), alignof(std::int64_t)); - container_starts_offset = reserve(containers, sizeof(std::int64_t), alignof(std::int64_t)); - num_containers_offset = reserve(1, sizeof(std::uint32_t), alignof(std::uint32_t)); - kinds_offset = reserve(containers, sizeof(container_kind), alignof(container_kind)); - payload_sizes_offset = reserve(containers, sizeof(std::uint64_t), alignof(std::uint64_t)); - payload_offsets_offset = reserve(containers, sizeof(std::uint64_t), alignof(std::uint64_t)); - summary_offset = reserve(1, sizeof(device_build_summary), alignof(device_build_summary)); - workspace_offset = reserve(workspace_bytes, sizeof(cuda::std::byte), alignof(std::max_align_t)); - bytes = cursor; - } - - std::size_t sorted_ids_offset{}; - std::size_t id_count_offset{}; - std::size_t valid_count_offset{}; - std::size_t selected_count_offset{}; - std::size_t container_starts_offset{}; - std::size_t num_containers_offset{}; - std::size_t kinds_offset{}; - std::size_t payload_sizes_offset{}; - std::size_t payload_offsets_offset{}; - std::size_t summary_offset{}; - std::size_t workspace_offset{}; - std::size_t bytes{}; -}; - -int grid_size_for(std::size_t count) -{ - auto const blocks = (count + kBuilderBlockSize - 1) / kBuilderBlockSize; - return static_cast(std::min(blocks, 65535)); -} - -__device__ void write_u16(cuda::std::byte* output, std::size_t offset, std::uint16_t value) -{ - auto* bytes = reinterpret_cast(output); - bytes[offset] = static_cast(value); - bytes[offset + 1] = static_cast(value >> 8); -} - -__device__ void write_u32(cuda::std::byte* output, std::size_t offset, std::uint32_t value) -{ - auto* bytes = reinterpret_cast(output); - for (int byte = 0; byte < 4; ++byte) { - bytes[offset + byte] = static_cast(value >> (8 * byte)); - } -} - -__device__ void write_u64(cuda::std::byte* output, std::size_t offset, std::uint64_t value) -{ - auto* bytes = reinterpret_cast(output); - for (int byte = 0; byte < 8; ++byte) { - bytes[offset + byte] = static_cast(value >> (8 * byte)); - } -} - -/** Find the sorted prefix that lies inside the logical dataset shape. */ -__global__ void find_valid_count_kernel(std::uint32_t const* ids, - std::int64_t const* id_count, - std::uint64_t dataset_rows, - std::int64_t* valid_count) -{ - if (blockIdx.x != 0 || threadIdx.x != 0) { return; } - std::int64_t first{}; - auto last = *id_count; - while (first < last) { - auto const middle = first + (last - first) / 2; - if (static_cast(ids[middle]) < dataset_rows) { - first = middle + 1; - } else { - last = middle; - } - } - *valid_count = first; -} - -/** Select the first sorted ID belonging to every high-16-bit container. */ -struct is_container_start { - std::uint32_t const* ids{}; - std::int64_t const* valid_count{}; - - __device__ bool operator()(std::int64_t i) const - { - auto const count = *valid_count; - return i < count && (i == 0 || (ids[i - 1] >> 16) != (ids[i] >> 16)); - } -}; - -__global__ void narrow_container_count_kernel(std::int64_t const* selected_count, - std::uint32_t* num_containers) -{ - if (blockIdx.x == 0 && threadIdx.x == 0) { - *num_containers = static_cast(*selected_count); - } -} - -/** - * Select the standard array or bitmap portable payload for each container. - * - * ID construction intentionally does not emit run containers. Full and nearly - * full allowlists should normally bypass filtering, and limiting construction - * to the two cuco-native forms keeps the batch builder and lookup behavior - * predictable. This still uses O(number of input IDs) scratch for sorting and - * scans; it never constructs a dense dataset-sized temporary bitmap. - */ -__global__ void analyze_containers_kernel(std::int64_t const* id_count, - std::int64_t const* container_starts, - std::uint32_t const* num_containers, - container_kind* kinds, - std::uint64_t* payload_sizes) -{ - auto container = - static_cast(static_cast(blockIdx.x) * blockDim.x + threadIdx.x); - auto const stride = static_cast(static_cast(gridDim.x) * - static_cast(blockDim.x)); - auto const count = *num_containers; - for (; container < count; container += stride) { - auto const begin = container_starts[container]; - auto const end = container + 1 < count ? container_starts[container + 1] : *id_count; - auto const cardinality = static_cast(end - begin); - if (cardinality <= kArrayCardinality) { - kinds[container] = container_kind::array; - payload_sizes[container] = cardinality * sizeof(std::uint16_t); - } else { - kinds[container] = container_kind::bitmap; - payload_sizes[container] = kBitmapBytes; - } - } -} - -/** Collect the scalar results needed to allocate the exact final portable byte - * stream. */ -__global__ void finish_device_analysis_kernel(std::int64_t const* id_count, - std::int64_t const* valid_count, - std::uint32_t const* num_containers, - std::uint64_t const* payload_sizes, - std::uint64_t const* payload_offsets, - device_build_summary* summary) -{ - if (blockIdx.x != 0 || threadIdx.x != 0) { return; } - auto const cardinality = *id_count; - auto const containers = *num_containers; - summary->cardinality = cardinality; - summary->num_containers = containers; - summary->payload_bytes = - containers == 0 ? 0 : payload_offsets[containers - 1] + payload_sizes[containers - 1]; - summary->invalid = cardinality != *valid_count; -} - -__host__ __device__ std::size_t portable_header_size(std::uint32_t num_containers) -{ - return 2 * sizeof(std::uint32_t) + - num_containers * (2 * sizeof(std::uint16_t) + sizeof(std::uint32_t)); -} - -/** Write the cookie, descriptors, and portable container-offset table. */ -__global__ void encode_header_kernel(std::uint32_t const* ids, - std::int64_t const* id_count, - std::int64_t const* container_starts, - std::uint32_t num_containers, - std::uint64_t const* payload_offsets, - std::size_t header_size, - cuda::std::byte* output) -{ - auto const thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - auto const stride = static_cast(gridDim.x) * blockDim.x; - constexpr auto descriptor_offset = 2 * sizeof(std::uint32_t); - auto const offsets_offset = descriptor_offset + num_containers * 2 * sizeof(std::uint16_t); - - if (thread == 0) { - write_u32(output, 0, kCookieNoRun); - write_u32(output, sizeof(std::uint32_t), num_containers); - } - - for (auto container = thread; container < num_containers; container += stride) { - auto const begin = container_starts[container]; - auto const end = container + 1 < num_containers ? container_starts[container + 1] : *id_count; - auto const descriptor = descriptor_offset + container * 2 * sizeof(std::uint16_t); - write_u16(output, descriptor, static_cast(ids[begin] >> 16)); - write_u16( - output, descriptor + sizeof(std::uint16_t), static_cast(end - begin - 1)); - write_u32(output, - offsets_offset + container * sizeof(std::uint32_t), - static_cast(header_size + payload_offsets[container])); - } -} - -/** Encode array and bitmap payloads directly into their final device offsets. */ -__global__ void encode_payloads_kernel(std::uint32_t const* ids, - std::int64_t const* id_count, - std::int64_t const* container_starts, - std::uint32_t num_containers, - container_kind const* kinds, - std::uint64_t const* payload_offsets, - std::size_t header_size, - cuda::std::byte* output) -{ - auto const container = static_cast(blockIdx.x); - if (container >= num_containers) { return; } - auto const begin = container_starts[container]; - auto const end = container + 1 < num_containers ? container_starts[container + 1] : *id_count; - auto const payload = header_size + payload_offsets[container]; - - if (kinds[container] == container_kind::array) { - for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { - write_u16(output, - payload + static_cast(i - begin) * sizeof(std::uint16_t), - static_cast(ids[i] & 0xffffu)); - } - return; - } - - __shared__ std::uint64_t bitmap_words[kBitmapBytes / sizeof(std::uint64_t)]; - constexpr std::uint32_t words = kBitmapBytes / sizeof(std::uint64_t); - for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { - bitmap_words[word] = 0; - } - __syncthreads(); - for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { - auto const lower = ids[i] & 0xffffu; - atomicOr(reinterpret_cast(&bitmap_words[lower / 64]), - static_cast(std::uint64_t{1} << (lower % 64))); - } - __syncthreads(); - for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { - write_u64(output, payload + word * sizeof(std::uint64_t), bitmap_words[word]); - } -} - -/** - * Sort and analyze an entire sparse allowlist in one CTA. - * - * The unsorted specialization uses a blocked 128 x 1 radix sort. Padding uses UINT32_MAX; - * writing only the first `size` sorted items is still correct when UINT32_MAX itself is a valid ID - * because all padding values compare equal to that final real value. Thread zero then walks at most - * 128 normalized IDs to - * build compact per-container metadata and the exact serialized payload size. - * - * Inputs are promised unique by the public API. This kernel deliberately does - * not spend work or storage checking or collapsing duplicates. - */ -template -__global__ void analyze_sparse_ids_kernel(std::uint32_t const* ids, - std::uint32_t size, - std::uint64_t dataset_rows, - std::uint32_t* sorted_ids, - sparse_container_metadata* metadata, - device_build_summary* summary) -{ - if constexpr (SortInput) { - using block_sort = - cub::BlockRadixSort; - __shared__ typename block_sort::TempStorage sort_storage; - std::uint32_t thread_ids[kSparseItemsPerThread]; - -#pragma unroll - for (int item = 0; item < kSparseItemsPerThread; ++item) { - auto const index = static_cast(threadIdx.x) * kSparseItemsPerThread + item; - thread_ids[item] = index < size ? ids[index] : std::numeric_limits::max(); - } - block_sort(sort_storage).Sort(thread_ids); -#pragma unroll - for (int item = 0; item < kSparseItemsPerThread; ++item) { - auto const index = static_cast(threadIdx.x) * kSparseItemsPerThread + item; - if (index < size) { sorted_ids[index] = thread_ids[item]; } - } - __syncthreads(); - } - - if (threadIdx.x != 0) { return; } - auto const* normalized_ids = SortInput ? sorted_ids : ids; - - summary->cardinality = size; - summary->payload_bytes = 0; - summary->num_containers = 0; - summary->invalid = 0; - - // Validate before writing container metadata. For a valid row, the logical - // dataset shape bounds the number of containers allocated by the host. - for (std::uint32_t i = 0; i < size; ++i) { - if (static_cast(normalized_ids[i]) >= dataset_rows) { - summary->invalid = 1; - return; - } - } - - std::uint32_t container{}; - std::uint32_t begin{}; - std::uint32_t payload_offset{}; - while (begin < size) { - auto const key = normalized_ids[begin] >> 16; - auto end = begin + 1; - while (end < size && (normalized_ids[end] >> 16) == key) { - ++end; - } - - auto const cardinality = end - begin; - auto const array_size = cardinality * sizeof(std::uint16_t); - metadata[container] = sparse_container_metadata{begin, payload_offset}; - payload_offset += array_size; - ++container; - begin = end; - } - - summary->payload_bytes = payload_offset; - summary->num_containers = container; -} - -/** - * Encode a sparse row in one CTA after exact output allocation. - * - * Thread zero writes every header byte, so this path needs no output memset. - * Array values are striped across the CTA. Bitmap payloads cannot occur below - * the sparse cardinality threshold. - */ -__global__ void encode_sparse_row_kernel(std::uint32_t const* ids, - std::uint32_t cardinality, - sparse_container_metadata const* metadata, - std::uint32_t num_containers, - std::size_t header_size, - cuda::std::byte* output, - ref_type* reference) -{ - constexpr auto descriptor_offset = 2 * sizeof(std::uint32_t); - auto const offsets_offset = descriptor_offset + num_containers * 2 * sizeof(std::uint16_t); - - if (threadIdx.x == 0) { - write_u32(output, 0, kCookieNoRun); - write_u32(output, sizeof(std::uint32_t), num_containers); - for (std::uint32_t container = 0; container < num_containers; ++container) { - auto const begin = metadata[container].begin; - auto const end = container + 1 < num_containers ? metadata[container + 1].begin : cardinality; - auto const descriptor = descriptor_offset + container * 2 * sizeof(std::uint16_t); - write_u16(output, descriptor, static_cast(ids[begin] >> 16)); - write_u16( - output, descriptor + sizeof(std::uint16_t), static_cast(end - begin - 1)); - write_u32(output, - offsets_offset + container * sizeof(std::uint32_t), - static_cast(header_size + metadata[container].payload_offset)); - } - } - - for (std::uint32_t container = 0; container < num_containers; ++container) { - auto const begin = metadata[container].begin; - auto const end = container + 1 < num_containers ? metadata[container + 1].begin : cardinality; - auto const payload = header_size + metadata[container].payload_offset; - for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { - write_u16(output, - payload + static_cast(i - begin) * sizeof(std::uint16_t), - static_cast(ids[i] & 0xffffu)); - } - } - __syncthreads(); - if (threadIdx.x == 0) { ::new (static_cast(reference)) ref_type{output}; } + auto const i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < size && static_cast(ids[i]) >= dataset_rows) { atomicExch(invalid, 1u); } } -device_build_result build_sparse_from_device_ids( - raft::resources const& res, - std::size_t dataset_rows, - raft::device_vector_view ids, - bool pre_sorted) +__global__ void store_ref_kernel(ref_type ref, ref_type* output) { - common::nvtx::range build_scope("roaring_allowlist::build_sparse"); - auto const stream = raft::resource::get_cuda_stream(res); - auto const size = static_cast(ids.extent(0)); - auto const num_chunks = - (static_cast(dataset_rows) + (std::uint64_t{1} << 16) - 1) >> 16; - auto const max_containers = std::min(size, static_cast(num_chunks)); - sparse_scratch_layout const layout{size, max_containers, !pre_sorted}; - rmm::device_uvector scratch(layout.bytes, stream); - - auto* summary = reinterpret_cast(scratch.data()); - auto* sorted_ids = - pre_sorted ? nullptr - : reinterpret_cast(scratch.data() + layout.sorted_ids_offset); - auto* metadata = - reinterpret_cast(scratch.data() + layout.metadata_offset); - - { - common::nvtx::range stage_scope( - "roaring_allowlist::sparse_analysis"); - if (pre_sorted) { - analyze_sparse_ids_kernel - <<<1, kSparseBuilderBlockSize, 0, stream>>>(ids.data_handle(), - static_cast(size), - dataset_rows, - nullptr, - metadata, - summary); - } else { - analyze_sparse_ids_kernel - <<<1, kSparseBuilderBlockSize, 0, stream>>>(ids.data_handle(), - static_cast(size), - dataset_rows, - sorted_ids, - metadata, - summary); - } - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } - - device_build_summary host_summary; - { - common::nvtx::range stage_scope("roaring_allowlist::size_readback"); - RAFT_CUDA_TRY(cudaMemcpyAsync( - &host_summary, summary, sizeof(host_summary), cudaMemcpyDeviceToHost, stream)); - raft::resource::sync_stream(res); - } - RAFT_EXPECTS(host_summary.invalid == 0, - "Roaring allowlist ID must be smaller than dataset_rows."); - RAFT_EXPECTS(host_summary.cardinality > 0 && host_summary.num_containers > 0, - "Internal error: nonempty sparse Roaring input produced an empty device build."); - - auto const header_size = portable_header_size(host_summary.num_containers); - auto const serialized_bytes = header_size + static_cast(host_summary.payload_bytes); - rmm::device_uvector output(0, stream); - { - common::nvtx::range stage_scope( - "roaring_allowlist::final_allocation"); - output.resize(owned_storage_bytes(serialized_bytes), stream); - } - auto const* normalized_ids = pre_sorted ? ids.data_handle() : sorted_ids; - common::nvtx::range encode_scope( - "roaring_allowlist::sparse_encode_and_ref"); - encode_sparse_row_kernel<<<1, kSparseBuilderBlockSize, 0, stream>>>( - normalized_ids, - static_cast(size), - metadata, - host_summary.num_containers, - header_size, - output.data(), - reinterpret_cast(output.data() + reference_offset(serialized_bytes))); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - return {std::move(output), serialized_bytes, size, true}; + if (threadIdx.x == 0) { ::new (static_cast(output)) ref_type{ref}; } } -#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER -__global__ void validate_cuco_input_kernel(std::uint32_t const* ids, - std::size_t size, - std::uint64_t dataset_rows, - std::uint32_t* invalid) +__global__ void contains_kernel(ref_type const* reference, + bool empty, + std::uint64_t dataset_rows, + 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 && static_cast(ids[i]) >= dataset_rows) { atomicExch(invalid, 1u); } + if (i >= size) { return; } + auto const row = row_ids[i]; + output[i] = !empty && static_cast(row) < dataset_rows && reference->contains(row); } -__global__ void store_cuco_ref_kernel(ref_type ref, ref_type* output) -{ - if (threadIdx.x == 0) { ::new (static_cast(output)) ref_type{ref}; } -} -#endif +struct device_build_result { + std::unique_ptr owner; + rmm::device_uvector reference; + std::size_t cardinality{}; +}; -/** - * Build a standard portable Roaring row from device IDs. - * - * When cuCollections provides its raw-index factories, the one-row path delegates construction to - * `roaring_bitmap::from_indices` or `from_sorted_unique_indices`, retains that owner without - * copying its serialized payload, and initializes the cuVS device reference once. The code below - * those factories remains the compatibility implementation for the currently pinned cuco revision. - * Multi-row inputs never enter this function; they use the segmented cuVS builder. Neither path - * allocates storage proportional to `dataset_rows` bits. - * - * Pre-sorted ordering and uniqueness are unchecked caller promises. - */ device_build_result build_from_device_ids( raft::resources const& res, std::size_t dataset_rows, @@ -742,1107 +112,58 @@ device_build_result build_from_device_ids( 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 {rmm::device_uvector(0, stream), 0, 0, false}; } -#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER - { - // PR #839 owns the serialized bytes and already caches parsed metadata in - // its host-side ref. Keep that allocation alive and materialize only the - // lightweight ref in cuVS device storage; the payload is never recopied. - rmm::device_uvector invalid(1, stream); - RAFT_CUDA_TRY(cudaMemsetAsync(invalid.data(), 0, sizeof(std::uint32_t), stream)); - validate_cuco_input_kernel<<>>( - ids.data_handle(), size, dataset_rows, invalid.data()); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - std::uint32_t host_invalid{}; - RAFT_CUDA_TRY(cudaMemcpyAsync( - &host_invalid, invalid.data(), sizeof(host_invalid), cudaMemcpyDeviceToHost, stream)); - - 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); - // Both PR #839 factories perform their exact-size readback after all prior - // stream work, so the validation result is ready without another sync. - RAFT_EXPECTS(host_invalid == 0, "Roaring allowlist ID must be smaller than dataset_rows."); - - auto owner = std::make_unique(std::move(bitmap)); - auto const serialized_bytes = static_cast(owner->size_bytes()); - auto const cardinality = static_cast(owner->size()); - rmm::device_uvector output(sizeof(ref_type), stream); - store_cuco_ref_kernel<<<1, 1, 0, stream>>>(owner->ref(), - reinterpret_cast(output.data())); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - return {std::move(output), serialized_bytes, cardinality, true, std::move(owner)}; - } -#endif - auto const sparse_cutoff = pre_sorted ? kSparseBuilderMaxPreSortedIds : kSparseBuilderMaxIds; - if (size <= sparse_cutoff) { - return build_sparse_from_device_ids(res, dataset_rows, ids, pre_sorted); - } + if (size == 0) { return {nullptr, rmm::device_uvector(0, stream), 0}; } - auto const num_chunks = - (static_cast(dataset_rows) + (std::uint64_t{1} << 16) - 1) >> 16; - auto const max_containers = std::min(size, static_cast(num_chunks)); - auto const item_count = static_cast(size); - auto const container_slots = static_cast(max_containers); - constexpr int sort_end_bit = std::numeric_limits::digits; - - std::size_t sort_workspace_bytes{}; - std::size_t select_workspace_bytes{}; - std::size_t payload_scan_workspace_bytes{}; - if (!pre_sorted) { - RAFT_CUDA_TRY(cub::DeviceRadixSort::SortKeys(nullptr, - sort_workspace_bytes, - ids.data_handle(), - static_cast(nullptr), - item_count, - 0, - sort_end_bit, - stream)); - } - auto counting = thrust::make_counting_iterator(0); - RAFT_CUDA_TRY(cub::DeviceSelect::If(nullptr, - select_workspace_bytes, - counting, - static_cast(nullptr), - static_cast(nullptr), - item_count, - is_container_start{ids.data_handle(), nullptr}, - stream)); - RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, - payload_scan_workspace_bytes, - static_cast(nullptr), - static_cast(nullptr), - container_slots, - stream)); - auto const workspace_bytes = - std::max(sort_workspace_bytes, std::max(select_workspace_bytes, payload_scan_workspace_bytes)); - general_scratch_layout const layout{size, max_containers, !pre_sorted, workspace_bytes}; - rmm::device_uvector scratch(layout.bytes, stream); - - auto* sorted_ids = - pre_sorted ? nullptr - : reinterpret_cast(scratch.data() + layout.sorted_ids_offset); - auto* id_count = reinterpret_cast(scratch.data() + layout.id_count_offset); - auto* valid_count = reinterpret_cast(scratch.data() + layout.valid_count_offset); - auto* selected_count = - reinterpret_cast(scratch.data() + layout.selected_count_offset); - auto* container_starts = - reinterpret_cast(scratch.data() + layout.container_starts_offset); - auto* num_containers = - reinterpret_cast(scratch.data() + layout.num_containers_offset); - auto* kinds = reinterpret_cast(scratch.data() + layout.kinds_offset); - auto* payload_sizes = - reinterpret_cast(scratch.data() + layout.payload_sizes_offset); - auto* payload_offsets = - reinterpret_cast(scratch.data() + layout.payload_offsets_offset); - auto* device_summary = - reinterpret_cast(scratch.data() + layout.summary_offset); - auto* workspace = scratch.data() + layout.workspace_offset; - - if (!pre_sorted) { - common::nvtx::range stage_scope("roaring_allowlist::radix_sort"); - RAFT_CUDA_TRY(cub::DeviceRadixSort::SortKeys(workspace, - sort_workspace_bytes, - ids.data_handle(), - sorted_ids, - item_count, - 0, - sort_end_bit, - stream)); - } - auto const* normalized_ids = pre_sorted ? ids.data_handle() : sorted_ids; - RAFT_CUDA_TRY( - cudaMemcpyAsync(id_count, &item_count, sizeof(item_count), cudaMemcpyHostToDevice, stream)); - - { - common::nvtx::range stage_scope( - "roaring_allowlist::container_discovery"); - find_valid_count_kernel<<<1, 1, 0, stream>>>( - normalized_ids, id_count, dataset_rows, valid_count); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - RAFT_CUDA_TRY(cub::DeviceSelect::If(workspace, - select_workspace_bytes, - counting, - container_starts, - selected_count, - item_count, - is_container_start{normalized_ids, valid_count}, - stream)); - narrow_container_count_kernel<<<1, 1, 0, stream>>>(selected_count, num_containers); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } - - { - common::nvtx::range stage_scope( - "roaring_allowlist::container_analysis"); - RAFT_CUDA_TRY( - cudaMemsetAsync(payload_sizes, 0, max_containers * sizeof(std::uint64_t), stream)); - analyze_containers_kernel<<>>( - valid_count, container_starts, num_containers, kinds, payload_sizes); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(workspace, - payload_scan_workspace_bytes, - payload_sizes, - payload_offsets, - container_slots, - stream)); - finish_device_analysis_kernel<<<1, 1, 0, stream>>>( - id_count, valid_count, num_containers, payload_sizes, payload_offsets, device_summary); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } - - device_build_summary summary; - { - common::nvtx::range stage_scope("roaring_allowlist::size_readback"); - RAFT_CUDA_TRY( - cudaMemcpyAsync(&summary, device_summary, sizeof(summary), cudaMemcpyDeviceToHost, stream)); - raft::resource::sync_stream(res); - } - RAFT_EXPECTS(summary.invalid == 0, "Roaring allowlist ID must be smaller than dataset_rows."); - RAFT_EXPECTS(summary.cardinality > 0 && summary.num_containers > 0, - "Internal error: nonempty Roaring input produced an empty device build."); - - auto const header_size = portable_header_size(summary.num_containers); - RAFT_EXPECTS(summary.payload_bytes <= std::numeric_limits::max() - header_size, - "Portable Roaring row exceeds the 32-bit offset range."); - auto const serialized_bytes = header_size + static_cast(summary.payload_bytes); - rmm::device_uvector output(0, stream); - { - common::nvtx::range stage_scope( - "roaring_allowlist::final_allocation"); - output.resize(owned_storage_bytes(serialized_bytes), stream); - } - auto const header_items = static_cast(summary.num_containers); - { - common::nvtx::range stage_scope("roaring_allowlist::header_encode"); - encode_header_kernel<<>>( - normalized_ids, - valid_count, - container_starts, - summary.num_containers, - payload_offsets, - header_size, - output.data()); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } - common::nvtx::range payload_scope( - "roaring_allowlist::payload_encode"); - encode_payloads_kernel<<>>( - normalized_ids, - valid_count, - container_starts, - summary.num_containers, - kinds, - payload_offsets, - header_size, - output.data()); + rmm::device_uvector invalid(1, stream); + RAFT_CUDA_TRY(cudaMemsetAsync(invalid.data(), 0, sizeof(std::uint32_t), stream)); + validate_input_kernel<<>>( + ids.data_handle(), size, dataset_rows, invalid.data()); RAFT_CUDA_TRY(cudaPeekAtLastError()); - return { - std::move(output), serialized_bytes, static_cast(summary.cardinality), false}; -} - -struct batched_device_build_result { - rmm::device_uvector storage; - std::vector row_offsets; - std::vector reference_offsets; - std::vector serialized_bytes; - std::vector cardinalities; - bool references_initialized{}; -#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER - std::unique_ptr cuco_owner{}; -#endif -}; - -struct packed_rows_layout { - std::vector row_offsets; - std::vector reference_offsets; - std::size_t references_offset{}; - std::size_t bytes{}; -}; -packed_rows_layout make_packed_rows_layout(std::vector const& serialized_bytes) -{ - packed_rows_layout result; - result.row_offsets.resize(serialized_bytes.size()); - result.reference_offsets.resize(serialized_bytes.size()); - std::size_t cursor{}; - bool any_nonempty{}; - for (std::size_t row = 0; row < serialized_bytes.size(); ++row) { - if (serialized_bytes[row] == 0) { continue; } - any_nonempty = true; - cursor = align_up(cursor, alignof(std::max_align_t)); - result.row_offsets[row] = cursor; - cursor += serialized_bytes[row]; - } - if (!any_nonempty) { return result; } - result.references_offset = align_up(cursor, alignof(ref_type)); - for (std::size_t row = 0; row < serialized_bytes.size(); ++row) { - result.reference_offsets[row] = result.references_offset + row * sizeof(ref_type); - } - result.bytes = result.references_offset + serialized_bytes.size() * sizeof(ref_type); - return result; -} - -struct batch_general_scratch_layout { - batch_general_scratch_layout(std::size_t total_ids, - std::size_t rows, - std::size_t max_containers, - bool store_sorted_ids, - std::size_t workspace_bytes) - { - std::size_t cursor{}; - auto reserve = [&](std::size_t count, std::size_t item_size, std::size_t alignment) { - auto const result = align_up(cursor, alignment); - cursor = result + count * item_size; - return result; - }; - if (store_sorted_ids) { - sorted_ids_offset = reserve(total_ids, sizeof(std::uint32_t), alignof(std::uint32_t)); - } - valid_counts_offset = reserve(rows, sizeof(std::int64_t), alignof(std::int64_t)); - flags_offset = reserve(total_ids, sizeof(std::uint8_t), alignof(std::uint8_t)); - selected_count_offset = reserve(1, sizeof(std::int64_t), alignof(std::int64_t)); - container_starts_offset = reserve(max_containers, sizeof(std::int64_t), alignof(std::int64_t)); - row_container_offsets_offset = reserve(rows + 1, sizeof(std::int64_t), alignof(std::int64_t)); - container_rows_offset = reserve(max_containers, sizeof(std::uint32_t), alignof(std::uint32_t)); - kinds_offset = reserve(max_containers, sizeof(container_kind), alignof(container_kind)); - payload_sizes_offset = reserve(max_containers, sizeof(std::uint64_t), alignof(std::uint64_t)); - payload_offsets_offset = reserve(max_containers, sizeof(std::uint64_t), alignof(std::uint64_t)); - summaries_offset = reserve(rows, sizeof(device_build_summary), alignof(device_build_summary)); - output_offsets_offset = reserve(rows, sizeof(std::uint64_t), alignof(std::uint64_t)); - workspace_offset = reserve(workspace_bytes, sizeof(cuda::std::byte), alignof(std::max_align_t)); - bytes = cursor; - } - - std::size_t sorted_ids_offset{}; - std::size_t valid_counts_offset{}; - std::size_t flags_offset{}; - std::size_t selected_count_offset{}; - std::size_t container_starts_offset{}; - std::size_t row_container_offsets_offset{}; - std::size_t container_rows_offset{}; - std::size_t kinds_offset{}; - std::size_t payload_sizes_offset{}; - std::size_t payload_offsets_offset{}; - std::size_t summaries_offset{}; - std::size_t output_offsets_offset{}; - std::size_t workspace_offset{}; - std::size_t bytes{}; -}; - -__global__ void find_batch_valid_counts_kernel(std::uint32_t const* ids, - std::int64_t const* indptr, - std::int64_t rows, - std::uint64_t dataset_rows, - std::int64_t* valid_counts) -{ - auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (row >= rows) { return; } - auto const begin = indptr[row]; - auto const width = indptr[row + 1] - begin; - std::int64_t first{}; - auto last = width; - while (first < last) { - auto const middle = first + (last - first) / 2; - if (static_cast(ids[begin + middle]) < dataset_rows) { - first = middle + 1; - } else { - last = middle; - } - } - valid_counts[row] = first; -} - -__global__ void mark_batch_container_starts_kernel(std::uint32_t const* ids, - std::int64_t const* indptr, - std::int64_t const* valid_counts, - std::int64_t rows, - std::uint8_t* flags) -{ - auto const row = static_cast(blockIdx.x); - if (row >= rows) { return; } - auto const begin = indptr[row]; - auto const width = indptr[row + 1] - begin; - auto const valid = valid_counts[row]; - for (std::int64_t column = threadIdx.x; column < width; column += blockDim.x) { - auto const index = begin + column; - flags[index] = - column < valid && (column == 0 || (ids[index - 1] >> 16) != (ids[index] >> 16)) ? 1 : 0; - } -} - -__global__ void find_row_container_offsets_kernel(std::int64_t const* container_starts, - std::int64_t const* selected_count, - std::int64_t const* indptr, - std::int64_t rows, - std::int64_t* row_offsets) -{ - auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (row > rows) { return; } - auto const target = indptr[row]; - std::int64_t first{}; - auto last = *selected_count; - while (first < last) { - auto const middle = first + (last - first) / 2; - if (container_starts[middle] < target) { - first = middle + 1; - } else { - last = middle; - } - } - row_offsets[row] = first; -} - -__global__ void fill_container_rows_kernel(std::int64_t const* row_container_offsets, - std::int64_t rows, - std::uint32_t* container_rows) -{ - auto const row = static_cast(blockIdx.x); - if (row >= rows) { return; } - for (auto container = row_container_offsets[row] + threadIdx.x; - container < row_container_offsets[row + 1]; - container += blockDim.x) { - container_rows[container] = static_cast(row); - } -} - -__global__ void analyze_batch_containers_kernel(std::int64_t const* indptr, - std::int64_t const* valid_counts, - std::int64_t const* container_starts, - std::int64_t const* selected_count, - std::uint32_t const* container_rows, - container_kind* kinds, - std::uint64_t* payload_sizes) -{ - auto container = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - auto const stride = static_cast(gridDim.x) * blockDim.x; - auto const count = *selected_count; - for (; container < count; container += stride) { - auto const row = static_cast(container_rows[container]); - auto const begin = container_starts[container]; - auto const row_end = indptr[row] + valid_counts[row]; - auto const end = container + 1 < count && container_rows[container + 1] == row - ? container_starts[container + 1] - : row_end; - auto const cardinality = static_cast(end - begin); - if (cardinality <= kArrayCardinality) { - kinds[container] = container_kind::array; - payload_sizes[container] = cardinality * sizeof(std::uint16_t); - } else { - kinds[container] = container_kind::bitmap; - payload_sizes[container] = kBitmapBytes; - } - } -} - -__global__ void finish_batch_rows_kernel(std::int64_t rows, - std::int64_t const* indptr, - std::int64_t const* valid_counts, - std::int64_t const* row_container_offsets, - std::uint64_t const* payload_sizes, - std::uint64_t const* payload_offsets, - device_build_summary* summaries) -{ - auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (row >= rows) { return; } - auto const begin = row_container_offsets[row]; - auto const end = row_container_offsets[row + 1]; - auto const cardinality = indptr[row + 1] - indptr[row]; - auto& summary = summaries[row]; - summary.cardinality = cardinality; - summary.num_containers = static_cast(end - begin); - summary.payload_bytes = - begin == end ? 0 : payload_offsets[end - 1] + payload_sizes[end - 1] - payload_offsets[begin]; - summary.invalid = valid_counts[row] != cardinality; -} - -__global__ void encode_batch_headers_kernel(std::uint32_t const* ids, - std::int64_t const* indptr, - std::int64_t const* valid_counts, - std::int64_t const* container_starts, - std::int64_t const* row_container_offsets, - std::uint64_t const* payload_offsets, - device_build_summary const* summaries, - std::uint64_t const* output_offsets, - cuda::std::byte* storage) -{ - auto const row = static_cast(blockIdx.x); - auto const summary = summaries[row]; - if (summary.num_containers == 0) { return; } - auto const first_container = row_container_offsets[row]; - auto* output = storage + output_offsets[row]; - constexpr auto descriptor_offset = 2 * sizeof(std::uint32_t); - auto const offsets_offset = - descriptor_offset + summary.num_containers * 2 * sizeof(std::uint16_t); - auto const header_size = portable_header_size(summary.num_containers); - if (threadIdx.x == 0) { - write_u32(output, 0, kCookieNoRun); - write_u32(output, sizeof(std::uint32_t), summary.num_containers); - } - auto const row_end = indptr[row] + valid_counts[row]; - for (std::uint32_t local = threadIdx.x; local < summary.num_containers; local += blockDim.x) { - auto const container = first_container + local; - auto const begin = container_starts[container]; - auto const end = local + 1 < summary.num_containers ? container_starts[container + 1] : row_end; - auto const descriptor = descriptor_offset + local * 2 * sizeof(std::uint16_t); - write_u16(output, descriptor, static_cast(ids[begin] >> 16)); - write_u16( - output, descriptor + sizeof(std::uint16_t), static_cast(end - begin - 1)); - auto const local_payload = payload_offsets[container] - payload_offsets[first_container]; - write_u32(output, - offsets_offset + local * sizeof(std::uint32_t), - static_cast(header_size + local_payload)); - } -} - -__global__ void encode_batch_payloads_kernel(std::uint32_t const* ids, - std::int64_t const* indptr, - std::int64_t const* valid_counts, - std::int64_t const* container_starts, - std::int64_t const* row_container_offsets, - std::uint32_t const* container_rows, - std::int64_t num_containers, - container_kind const* kinds, - std::uint64_t const* payload_offsets, - device_build_summary const* summaries, - std::uint64_t const* output_offsets, - cuda::std::byte* storage) -{ - auto const container = static_cast(blockIdx.x); - if (container >= num_containers) { return; } - auto const row = static_cast(container_rows[container]); - auto const begin = container_starts[container]; - auto const first_container = row_container_offsets[row]; - auto const local = container - first_container; - auto const summary = summaries[row]; - auto const row_end = indptr[row] + valid_counts[row]; - auto const end = local + 1 < summary.num_containers ? container_starts[container + 1] : row_end; - auto* output = storage + output_offsets[row]; - auto const payload = portable_header_size(summary.num_containers) + payload_offsets[container] - - payload_offsets[first_container]; - - if (kinds[container] == container_kind::array) { - for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { - write_u16(output, - payload + static_cast(i - begin) * sizeof(std::uint16_t), - static_cast(ids[i] & 0xffffu)); - } - return; - } - - __shared__ std::uint64_t bitmap_words[kBitmapBytes / sizeof(std::uint64_t)]; - constexpr std::uint32_t words = kBitmapBytes / sizeof(std::uint64_t); - for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { - bitmap_words[word] = 0; - } - __syncthreads(); - for (auto i = begin + threadIdx.x; i < end; i += blockDim.x) { - auto const lower = ids[i] & 0xffffu; - atomicOr(reinterpret_cast(&bitmap_words[lower / 64]), - static_cast(std::uint64_t{1} << (lower % 64))); - } - __syncthreads(); - for (std::uint32_t word = threadIdx.x; word < words; word += blockDim.x) { - write_u64(output, payload + word * sizeof(std::uint64_t), bitmap_words[word]); - } -} - -__global__ void initialize_batch_refs_kernel(cuda::std::byte const* storage, - std::uint64_t const* row_offsets, - device_build_summary const* summaries, - ref_type* references, - std::size_t rows) -{ - auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (row < rows && summaries[row].num_containers != 0) { - ::new (static_cast(references + row)) ref_type{storage + row_offsets[row]}; - } -} - -batched_device_build_result build_general_rows( - raft::resources const& res, - std::size_t dataset_rows, - raft::device_vector_view ids, - raft::device_vector_view indptr, - std::vector const& host_indptr, - bool pre_sorted) -{ - common::nvtx::range build_scope( - "roaring_allowlist::build_general_batch"); - auto const stream = raft::resource::get_cuda_stream(res); - auto const rows = host_indptr.size() - 1; - auto const total_ids = static_cast(ids.extent(0)); - auto const item_count = static_cast(total_ids); - auto const row_count = static_cast(rows); - auto const chunks = - (static_cast(dataset_rows) + (std::uint64_t{1} << 16) - 1) >> 16; - std::size_t max_containers{}; - for (std::size_t row = 0; row < rows; ++row) { - auto const width = static_cast(host_indptr[row + 1] - host_indptr[row]); - max_containers += std::min(width, static_cast(chunks)); - } - RAFT_EXPECTS(rows <= std::numeric_limits::max(), - "Batched Roaring construction has too many rows for one launch."); - RAFT_EXPECTS(max_containers <= std::numeric_limits::max(), - "Batched Roaring construction has too many containers for one launch."); - constexpr int sort_end_bit = std::numeric_limits::digits; - auto counting = thrust::make_counting_iterator(0); - - std::size_t sort_workspace_bytes{}; - std::size_t select_workspace_bytes{}; - std::size_t scan_workspace_bytes{}; - if (!pre_sorted) { - RAFT_CUDA_TRY(cub::DeviceSegmentedRadixSort::SortKeys(nullptr, - sort_workspace_bytes, - ids.data_handle(), - static_cast(nullptr), - item_count, - row_count, - indptr.data_handle(), - indptr.data_handle() + 1, - 0, - sort_end_bit, - stream)); - } - RAFT_CUDA_TRY(cub::DeviceSelect::Flagged(nullptr, - select_workspace_bytes, - counting, - static_cast(nullptr), - static_cast(nullptr), - static_cast(nullptr), - item_count, - stream)); - RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, - scan_workspace_bytes, - static_cast(nullptr), - static_cast(nullptr), - static_cast(max_containers), - stream)); - auto const workspace_bytes = - std::max(sort_workspace_bytes, std::max(select_workspace_bytes, scan_workspace_bytes)); - batch_general_scratch_layout const layout{ - total_ids, rows, max_containers, !pre_sorted, workspace_bytes}; - rmm::device_uvector scratch(layout.bytes, stream); - auto* sorted_ids = - pre_sorted ? nullptr - : reinterpret_cast(scratch.data() + layout.sorted_ids_offset); - auto* valid_counts = reinterpret_cast(scratch.data() + layout.valid_counts_offset); - auto* flags = reinterpret_cast(scratch.data() + layout.flags_offset); - auto* selected_count = - reinterpret_cast(scratch.data() + layout.selected_count_offset); - auto* container_starts = - reinterpret_cast(scratch.data() + layout.container_starts_offset); - auto* row_container_offsets = - reinterpret_cast(scratch.data() + layout.row_container_offsets_offset); - auto* container_rows = - reinterpret_cast(scratch.data() + layout.container_rows_offset); - auto* kinds = reinterpret_cast(scratch.data() + layout.kinds_offset); - auto* payload_sizes = - reinterpret_cast(scratch.data() + layout.payload_sizes_offset); - auto* payload_offsets = - reinterpret_cast(scratch.data() + layout.payload_offsets_offset); - auto* summaries = - reinterpret_cast(scratch.data() + layout.summaries_offset); - auto* output_offsets = - reinterpret_cast(scratch.data() + layout.output_offsets_offset); - auto* workspace = scratch.data() + layout.workspace_offset; - - if (!pre_sorted) { - RAFT_CUDA_TRY(cub::DeviceSegmentedRadixSort::SortKeys(workspace, - sort_workspace_bytes, - ids.data_handle(), - sorted_ids, - item_count, - row_count, - indptr.data_handle(), - indptr.data_handle() + 1, - 0, - sort_end_bit, - stream)); - } - auto const* normalized = pre_sorted ? ids.data_handle() : sorted_ids; - find_batch_valid_counts_kernel<<>>( - normalized, indptr.data_handle(), row_count, dataset_rows, valid_counts); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - mark_batch_container_starts_kernel<<(rows), - kBuilderBlockSize, - 0, - stream>>>( - normalized, indptr.data_handle(), valid_counts, row_count, flags); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - RAFT_CUDA_TRY(cub::DeviceSelect::Flagged(workspace, - select_workspace_bytes, - counting, - flags, - container_starts, - selected_count, - item_count, - stream)); - find_row_container_offsets_kernel<<>>( - container_starts, selected_count, indptr.data_handle(), row_count, row_container_offsets); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - fill_container_rows_kernel<<(rows), kBuilderBlockSize, 0, stream>>>( - row_container_offsets, row_count, container_rows); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - RAFT_CUDA_TRY(cudaMemsetAsync(payload_sizes, 0, max_containers * sizeof(std::uint64_t), stream)); - analyze_batch_containers_kernel<<>>( - indptr.data_handle(), - valid_counts, - container_starts, - selected_count, - container_rows, - kinds, - payload_sizes); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(workspace, - scan_workspace_bytes, - payload_sizes, - payload_offsets, - static_cast(max_containers), - stream)); - finish_batch_rows_kernel<<>>( - row_count, - indptr.data_handle(), - valid_counts, - row_container_offsets, - payload_sizes, - payload_offsets, - summaries); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - - std::vector host_summaries(rows); - RAFT_CUDA_TRY(cudaMemcpyAsync(host_summaries.data(), - summaries, - rows * sizeof(device_build_summary), - cudaMemcpyDeviceToHost, - stream)); - raft::resource::sync_stream(res); - std::vector serialized(rows); - std::vector cardinalities(rows); - std::size_t actual_containers{}; - for (std::size_t row = 0; row < rows; ++row) { - auto const& summary = host_summaries[row]; - auto const cardinality = static_cast(host_indptr[row + 1] - host_indptr[row]); - RAFT_EXPECTS(summary.invalid == 0, "Roaring allowlist ID must be smaller than dataset_rows."); - RAFT_EXPECTS(summary.cardinality == static_cast(cardinality) && - (cardinality == 0 || summary.num_containers > 0), - "Internal error: general batched row analysis failed."); - cardinalities[row] = cardinality; - if (cardinality == 0) { continue; } - auto const header = portable_header_size(summary.num_containers); - RAFT_EXPECTS(summary.payload_bytes <= std::numeric_limits::max() - header, - "Portable Roaring row exceeds the 32-bit offset range."); - serialized[row] = header + static_cast(summary.payload_bytes); - actual_containers += summary.num_containers; - } - auto packed = make_packed_rows_layout(serialized); - rmm::device_uvector storage(packed.bytes, stream); - static_assert(sizeof(std::size_t) == sizeof(std::uint64_t)); - RAFT_CUDA_TRY(cudaMemcpyAsync(output_offsets, - packed.row_offsets.data(), - rows * sizeof(std::uint64_t), - cudaMemcpyHostToDevice, - stream)); - encode_batch_headers_kernel<<(rows), kBuilderBlockSize, 0, stream>>>( - normalized, - indptr.data_handle(), - valid_counts, - container_starts, - row_container_offsets, - payload_offsets, - summaries, - output_offsets, - storage.data()); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - encode_batch_payloads_kernel<<(actual_containers), - kBuilderBlockSize, - 0, - stream>>>(normalized, - indptr.data_handle(), - valid_counts, - container_starts, - row_container_offsets, - container_rows, - static_cast(actual_containers), - kinds, - payload_offsets, - summaries, - output_offsets, - storage.data()); + std::uint32_t host_invalid{}; + RAFT_CUDA_TRY(cudaMemcpyAsync( + &host_invalid, invalid.data(), sizeof(host_invalid), cudaMemcpyDeviceToHost, stream)); + + 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); + + // The exact-size readback in the cuco factory orders the preceding validation copy. + RAFT_EXPECTS(host_invalid == 0, "Roaring allowlist ID must be smaller than dataset_rows."); + + auto owner = std::make_unique(std::move(bitmap)); + auto const cardinality = static_cast(owner->size()); + rmm::device_uvector reference(sizeof(ref_type), stream); + store_ref_kernel<<<1, 1, 0, stream>>>(owner->ref(), + reinterpret_cast(reference.data())); RAFT_CUDA_TRY(cudaPeekAtLastError()); - initialize_batch_refs_kernel<<>>( - storage.data(), - output_offsets, - summaries, - reinterpret_cast(storage.data() + packed.references_offset), - rows); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - return {std::move(storage), - std::move(packed.row_offsets), - std::move(packed.reference_offsets), - std::move(serialized), - std::move(cardinalities), - true}; -} - -void validate_indptr(std::vector const& indptr, std::size_t rows, std::size_t nnz) -{ - RAFT_EXPECTS(rows > 0, "Roaring input must contain at least one allowlist row."); - RAFT_EXPECTS(indptr.size() == rows + 1, "Roaring indptr must contain num_rows + 1 entries."); - RAFT_EXPECTS(indptr.front() == 0, "Roaring indptr must start at zero."); - for (std::size_t row = 0; row < rows; ++row) { - RAFT_EXPECTS(indptr[row] <= indptr[row + 1] && indptr[row] >= 0, - "Roaring indptr must be nonnegative and nondecreasing."); - } - RAFT_EXPECTS(indptr.back() >= 0 && static_cast(indptr.back()) == nnz, - "The final Roaring indptr entry must equal nnz."); -} - -batched_device_build_result build_batched_from_device_ids( - raft::resources const& res, - std::size_t dataset_rows, - raft::device_vector_view ids, - raft::device_vector_view indptr, - std::vector const& host_indptr, - bool pre_sorted) -{ - auto const stream = raft::resource::get_cuda_stream(res); - auto const rows = host_indptr.size() - 1; - auto const size = static_cast(ids.extent(0)); - validate_indptr(host_indptr, rows, size); - if (size == 0) { - return {rmm::device_uvector(0, stream), - std::vector(rows), - std::vector(rows), - std::vector(rows), - std::vector(rows), - true}; - } - if (rows == 1) { - auto row_view = raft::make_device_vector_view( - ids.data_handle(), static_cast(size)); - auto built = build_from_device_ids(res, dataset_rows, row_view, pre_sorted); -#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER - return {std::move(built.storage), - {0}, - {0}, - {built.serialized_bytes}, - {built.cardinality}, - built.reference_initialized, - std::move(built.cuco_owner)}; -#else - return {std::move(built.storage), - {0}, - {built.cardinality == 0 ? 0 : reference_offset(built.serialized_bytes)}, - {built.serialized_bytes}, - {built.cardinality}, - built.reference_initialized}; -#endif - } - return build_general_rows(res, dataset_rows, ids, indptr, host_indptr, pre_sorted); -} - -/** - * Validate one externally supplied portable row before cuco sees it. - * - * cuco's raw-byte reference constructor assumes a valid stream and is not given - * the row's byte length. We therefore walk the complete row on the host first, - * checking every read, the required header variant, ordered keys and values, - * payload cardinalities, exact container offsets, the logical dataset bound, - * and that no trailing bytes remain. Besides rejecting malformed input, this - * walk records cardinality without retaining another decoded representation. - * - * A zero-length input is a cuVS convenience for an empty allowlist. The - * standard serialized empty form (no-run cookie followed by N = 0) is accepted - * as well. - */ -row_metadata validate_serialized_row(std::byte const* data, - std::size_t size, - std::size_t dataset_rows) -{ - // Empty here means the outer cuVS byte offsets selected no portable bytes for - // this row. - if (size == 0) { return {}; } - - // Decode the cookie first because it determines both how N is stored and - // whether a run bitmap follows. Every subsequent read advances `cursor` - // through exactly one format field. - auto const cookie = read_u32(data, size, 0); - bool const has_run = (cookie & 0xffffu) == kCookieRun; - std::size_t num_containers{}; - std::size_t cursor = 4; - if (has_run) { - num_containers = (cookie >> 16) + 1; - } else { - RAFT_EXPECTS(cookie == kCookieNoRun, "Malformed portable Roaring bitmap: unsupported cookie."); - num_containers = read_u32(data, size, cursor); - cursor += 4; - } - RAFT_EXPECTS(num_containers <= (std::size_t{1} << 16), - "Malformed portable Roaring bitmap: too many containers."); - if (num_containers == 0) { - RAFT_EXPECTS(!has_run && cursor == size, - "Malformed portable Roaring bitmap: invalid empty representation."); - return {}; - } - - std::size_t run_bitmap_offset{}; - if (has_run) { - auto const run_bitmap_bytes = (num_containers + 7) / 8; - RAFT_EXPECTS(cursor <= size && size - cursor >= run_bitmap_bytes, - "Malformed portable Roaring bitmap: truncated run bitmap."); - run_bitmap_offset = cursor; - cursor += run_bitmap_bytes; - } - - // The descriptive header is common to both cookie forms. Reconstruct - // cardinality by adding one to its encoded value and require container keys - // to be strictly increasing. - std::vector keys(num_containers); - std::vector cards(num_containers); - for (std::size_t i = 0; i < num_containers; ++i) { - keys[i] = read_u16(data, size, cursor); - cards[i] = static_cast(read_u16(data, size, cursor + 2)) + 1; - cursor += 4; - if (i > 0) { - RAFT_EXPECTS(keys[i - 1] < keys[i], - "Malformed portable Roaring bitmap: container keys are not ordered."); - } - } - - // These are offsets INSIDE this portable row. The no-run form always stores - // them; the run form stores them only at the specification's four-container - // threshold. - auto const store_offsets = !has_run || num_containers >= kOffsetThreshold; - std::vector container_offsets; - if (store_offsets) { - container_offsets.resize(num_containers); - for (std::size_t i = 0; i < num_containers; ++i) { - container_offsets[i] = read_u32(data, size, cursor); - cursor += 4; - } - } - - // Validate payloads in descriptor order. Requiring every stored offset to - // equal `cursor` also rejects gaps, overlaps, and offsets that point into a - // header or a different container. - row_metadata metadata; - metadata.empty = false; - for (std::size_t i = 0; i < num_containers; ++i) { - if (store_offsets) { - RAFT_EXPECTS(container_offsets[i] == cursor, - "Malformed portable Roaring bitmap: invalid container offset."); - } - auto const is_run = - has_run && ((std::to_integer(data[run_bitmap_offset + i / 8]) >> (i % 8)) & 1u); - std::uint32_t lower_max{}; - if (is_run) { - auto const num_runs = read_u16(data, size, cursor); - cursor += 2; - std::uint32_t run_cardinality{}; - std::uint32_t previous_end{}; - for (std::size_t run_index = 0; run_index < num_runs; ++run_index) { - auto const start = static_cast(read_u16(data, size, cursor)); - auto const length = static_cast(read_u16(data, size, cursor + 2)); - cursor += 4; - auto const end = start + length; - RAFT_EXPECTS(end <= std::numeric_limits::max(), - "Malformed portable Roaring bitmap: run exceeds uint16 range."); - if (run_index > 0) { - RAFT_EXPECTS(start > previous_end, - "Malformed portable Roaring bitmap: runs overlap or are " - "unordered."); - } - previous_end = end; - lower_max = end; - run_cardinality += length + 1; - } - RAFT_EXPECTS(num_runs > 0 && run_cardinality == cards[i], - "Malformed portable Roaring bitmap: invalid run cardinality."); - } else if (cards[i] <= kArrayCardinality) { - std::uint32_t previous{}; - for (std::size_t j = 0; j < cards[i]; ++j) { - auto const value = static_cast(read_u16(data, size, cursor)); - cursor += 2; - if (j > 0) { - RAFT_EXPECTS(previous < value, - "Malformed portable Roaring bitmap: " - "array values are not ordered."); - } - previous = value; - lower_max = value; - } - } else { - RAFT_EXPECTS(cursor <= size && size - cursor >= kBitmapBytes, - "Malformed portable Roaring bitmap: truncated bitmap container."); - std::size_t popcount{}; - bool found_max = false; - for (std::size_t j = 0; j < kBitmapBytes; ++j) { - popcount += std::popcount(std::to_integer(data[cursor + j])); - } - for (std::size_t j = kBitmapBytes; j-- > 0 && !found_max;) { - auto const byte = std::to_integer(data[cursor + j]); - if (byte != 0) { - lower_max = static_cast(j * 8 + (7 - std::countl_zero(byte))); - found_max = true; - } - } - RAFT_EXPECTS(found_max && popcount == cards[i], - "Malformed portable Roaring bitmap: invalid bitmap cardinality."); - cursor += kBitmapBytes; - } - - auto const max_id = (static_cast(keys[i]) << 16) | lower_max; - RAFT_EXPECTS(static_cast(max_id) < dataset_rows, - "Portable Roaring bitmap contains an ID outside dataset_rows."); - metadata.cardinality += cards[i]; - metadata.max_id = max_id; - } - RAFT_EXPECTS(cursor == size, "Malformed portable Roaring bitmap: trailing or unconsumed bytes."); - return metadata; -} - -/** - * Construct the lightweight cuco reference once, outside the search path. - * - * The raw-byte constructor parses the portable header and stores small - * container-location metadata by value while retaining pointers into `data`. It - * neither allocates nor copies the serialized payload. Empty allowlists skip - * this kernel because the pinned parser expects at least one container. - */ -__global__ void initialize_imported_refs_kernel(cuda::std::byte const* storage, - std::uint64_t const* row_offsets, - std::uint64_t const* serialized_bytes, - ref_type* references, - std::size_t rows) -{ - auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (row < rows && serialized_bytes[row] != 0) { - ::new (static_cast(references + row)) ref_type{storage + row_offsets[row]}; - } -} - -__global__ void initialize_view_tables_kernel(cuda::std::byte const* storage, - std::uint64_t const* reference_offsets, - std::uint64_t const* serialized_bytes, - ref_type const** references, - std::uint8_t* empty_rows, - std::size_t rows) -{ - auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (row >= rows) { return; } - auto const empty = serialized_bytes[row] == 0; - references[row] = - empty ? nullptr : reinterpret_cast(storage + reference_offsets[row]); - empty_rows[row] = empty ? 1 : 0; -} - -__global__ void contains_kernel(ref_type const* const* references, - std::uint8_t const* empty_rows, - std::uint64_t dataset_rows, - std::uint32_t const* row_ids, - std::uint8_t* output, - std::size_t columns, - std::size_t size) -{ - auto const i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (i >= size) { return; } - auto const query = i / columns; - auto const row = row_ids[i]; - output[i] = empty_rows[query] == 0 && static_cast(row) < dataset_rows && - references[query]->contains(row); + return {std::move(owner), std::move(reference), cardinality}; } } // namespace struct roaring_allowlist::impl { -#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER - std::unique_ptr cuco_owner; -#endif - rmm::device_uvector storage; - std::vector row_offsets_; - std::vector reference_offsets_; - std::vector serialized_bytes_; - std::vector cardinalities_; - rmm::device_uvector references; - rmm::device_uvector empty_rows; + std::unique_ptr owner; + rmm::device_uvector reference_storage; + std::size_t cardinality_{}; std::size_t dataset_rows_{}; - std::size_t total_cardinality_{}; - bool references_initialized_{}; - impl(raft::resources const& res, std::size_t dataset_rows, batched_device_build_result&& built) -#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER - : cuco_owner(std::move(built.cuco_owner)), - storage(std::move(built.storage)), -#else - : storage(std::move(built.storage)), -#endif - row_offsets_(std::move(built.row_offsets)), - reference_offsets_(std::move(built.reference_offsets)), - serialized_bytes_(std::move(built.serialized_bytes)), - cardinalities_(std::move(built.cardinalities)), - references(cardinalities_.size(), raft::resource::get_cuda_stream(res)), - empty_rows(cardinalities_.size(), raft::resource::get_cuda_stream(res)), - dataset_rows_(dataset_rows), - total_cardinality_( - std::accumulate(cardinalities_.begin(), cardinalities_.end(), std::size_t{})), - references_initialized_(built.references_initialized) + impl(std::size_t dataset_rows, device_build_result&& built) + : owner(std::move(built.owner)), + reference_storage(std::move(built.reference)), + cardinality_(built.cardinality), + dataset_rows_(dataset_rows) { static_assert(std::is_trivially_destructible_v); - auto const rows = cardinalities_.size(); - RAFT_EXPECTS(rows > 0 && row_offsets_.size() == rows && reference_offsets_.size() == rows && - serialized_bytes_.size() == rows, - "Internal error: inconsistent batched Roaring metadata."); - if (rows == 0) { return; } - - auto const stream = raft::resource::get_cuda_stream(res); - static_assert(sizeof(std::size_t) == sizeof(std::uint64_t)); - rmm::device_uvector device_row_offsets(rows, stream); - rmm::device_uvector device_reference_offsets(rows, stream); - rmm::device_uvector device_serialized_bytes(rows, stream); - RAFT_CUDA_TRY(cudaMemcpyAsync(device_row_offsets.data(), - row_offsets_.data(), - rows * sizeof(std::uint64_t), - cudaMemcpyHostToDevice, - stream)); - RAFT_CUDA_TRY(cudaMemcpyAsync(device_reference_offsets.data(), - reference_offsets_.data(), - rows * sizeof(std::uint64_t), - cudaMemcpyHostToDevice, - stream)); - RAFT_CUDA_TRY(cudaMemcpyAsync(device_serialized_bytes.data(), - serialized_bytes_.data(), - rows * sizeof(std::uint64_t), - cudaMemcpyHostToDevice, - stream)); - if (!references_initialized_ && storage.size() != 0) { - initialize_imported_refs_kernel<<>>( - storage.data(), - device_row_offsets.data(), - device_serialized_bytes.data(), - reinterpret_cast(storage.data() + reference_offsets_.front()), - rows); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - } - initialize_view_tables_kernel<<>>( - storage.data(), - device_reference_offsets.data(), - device_serialized_bytes.data(), - references.data(), - empty_rows.data(), - rows); - RAFT_CUDA_TRY(cudaPeekAtLastError()); } - [[nodiscard]] ref_type const* reference(std::size_t allowlist_id) const noexcept + [[nodiscard]] ref_type const* reference() const noexcept { - return cardinalities_[allowlist_id] == 0 - ? nullptr - : reinterpret_cast(storage.data() + reference_offsets_[allowlist_id]); + return cardinality_ == 0 ? nullptr + : reinterpret_cast(reference_storage.data()); } }; @@ -1854,166 +175,61 @@ roaring_allowlist roaring_allowlist::from_ids( raft::resources const& res, std::size_t dataset_rows, raft::host_vector_view ids, - raft::host_vector_view indptr, bool pre_sorted) { validate_dataset_rows(dataset_rows); - RAFT_EXPECTS(indptr.extent(0) >= 2, "Roaring indptr must contain at least two entries."); - auto const rows = static_cast(indptr.extent(0) - 1); - auto const size = static_cast(ids.extent(0)); - std::vector host_indptr(indptr.data_handle(), - indptr.data_handle() + indptr.extent(0)); - validate_indptr(host_indptr, rows, size); - + auto const size = static_cast(ids.extent(0)); auto const stream = raft::resource::get_cuda_stream(res); rmm::device_uvector device_ids(size, stream); - rmm::device_uvector device_indptr(rows + 1, stream); if (size != 0) { RAFT_CUDA_TRY(cudaMemcpyAsync(device_ids.data(), ids.data_handle(), size * sizeof(key_type), cudaMemcpyHostToDevice, stream)); - RAFT_CUDA_TRY(cudaMemcpyAsync(device_indptr.data(), - host_indptr.data(), - (rows + 1) * sizeof(indptr_type), - cudaMemcpyHostToDevice, - stream)); } auto device_ids_view = raft::make_device_vector_view( device_ids.data(), static_cast(size)); - auto device_indptr_view = raft::make_device_vector_view( - device_indptr.data(), static_cast(rows + 1)); - auto built = build_batched_from_device_ids( - res, dataset_rows, device_ids_view, device_indptr_view, host_indptr, pre_sorted); - return roaring_allowlist{std::make_unique(res, dataset_rows, std::move(built))}; + auto built = build_from_device_ids(res, dataset_rows, device_ids_view, pre_sorted); + return roaring_allowlist{std::make_unique(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, - raft::device_vector_view indptr, bool pre_sorted) { validate_dataset_rows(dataset_rows); - RAFT_EXPECTS(indptr.extent(0) >= 2, "Roaring indptr must contain at least two entries."); - auto const rows = static_cast(indptr.extent(0) - 1); - auto const size = static_cast(ids.extent(0)); - auto const stream = raft::resource::get_cuda_stream(res); - std::vector host_indptr(rows + 1); - RAFT_CUDA_TRY(cudaMemcpyAsync(host_indptr.data(), - indptr.data_handle(), - (rows + 1) * sizeof(indptr_type), - cudaMemcpyDeviceToHost, - stream)); - raft::resource::sync_stream(res); - validate_indptr(host_indptr, rows, size); - auto built = - build_batched_from_device_ids(res, dataset_rows, ids, indptr, host_indptr, pre_sorted); - return roaring_allowlist{std::make_unique(res, dataset_rows, std::move(built))}; -} - -roaring_allowlist roaring_allowlist::from_serialized( - raft::resources const& res, - std::size_t dataset_rows, - raft::host_vector_view bytes, - raft::host_vector_view byte_offsets) -{ - validate_dataset_rows(dataset_rows); - RAFT_EXPECTS(byte_offsets.extent(0) >= 2, - "Roaring byte_offsets must contain at least two entries."); - auto const rows = static_cast(byte_offsets.extent(0) - 1); - auto const size = static_cast(bytes.extent(0)); - RAFT_EXPECTS(byte_offsets(0) == 0, "Roaring byte_offsets must start at zero."); - RAFT_EXPECTS(byte_offsets(static_cast(rows)) == size, - "The final Roaring byte offset must equal bytes.extent(0)."); - - std::vector serialized_bytes(rows); - std::vector cardinalities(rows); - for (std::size_t row = 0; row < rows; ++row) { - auto const begin = byte_offsets(static_cast(row)); - auto const end = byte_offsets(static_cast(row + 1)); - RAFT_EXPECTS(begin <= end && end <= size, - "Roaring byte_offsets must be nondecreasing and in bounds."); - auto const row_size = static_cast(end - begin); - auto const metadata = - validate_serialized_row(bytes.data_handle() + begin, row_size, dataset_rows); - serialized_bytes[row] = metadata.empty ? 0 : row_size; - cardinalities[row] = metadata.cardinality; - } - - auto const packed = make_packed_rows_layout(serialized_bytes); - auto const stream = raft::resource::get_cuda_stream(res); - rmm::device_uvector storage(packed.bytes, stream); - for (std::size_t row = 0; row < rows; ++row) { - if (serialized_bytes[row] == 0) { continue; } - auto const begin = byte_offsets(static_cast(row)); - RAFT_CUDA_TRY(cudaMemcpyAsync(storage.data() + packed.row_offsets[row], - bytes.data_handle() + begin, - serialized_bytes[row], - cudaMemcpyHostToDevice, - stream)); - } - - batched_device_build_result built{std::move(storage), - packed.row_offsets, - packed.reference_offsets, - std::move(serialized_bytes), - std::move(cardinalities), - false}; - return roaring_allowlist{std::make_unique(res, dataset_rows, std::move(built))}; + auto built = build_from_device_ids(res, dataset_rows, ids, pre_sorted); + return roaring_allowlist{std::make_unique(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::num_allowlists() const noexcept -{ - return impl_->cardinalities_.size(); -} - std::size_t roaring_allowlist::dataset_rows() const noexcept { return impl_->dataset_rows_; } -std::size_t roaring_allowlist::cardinality(std::size_t allowlist_id) const -{ - RAFT_EXPECTS(allowlist_id < num_allowlists(), "Roaring allowlist_id is out of range."); - return impl_->cardinalities_[allowlist_id]; -} +std::size_t roaring_allowlist::cardinality() const noexcept { return impl_->cardinality_; } -bool roaring_allowlist::empty(std::size_t allowlist_id) const -{ - return cardinality(allowlist_id) == 0; -} - -std::size_t roaring_allowlist::total_cardinality() const noexcept -{ - return impl_->total_cardinality_; -} +bool roaring_allowlist::empty() const noexcept { return cardinality() == 0; } std::size_t roaring_allowlist::size_bytes() const noexcept { - auto bytes = impl_->storage.size() * sizeof(cuda::std::byte) + - impl_->references.size() * sizeof(ref_type const*) + - impl_->empty_rows.size() * sizeof(std::uint8_t); -#if CUVS_HAS_CUCO_ROARING_BITMAP_BUILDER - if (impl_->cuco_owner) { bytes += static_cast(impl_->cuco_owner->size_bytes()); } -#endif + auto bytes = impl_->reference_storage.size() * sizeof(cuda::std::byte); + if (impl_->owner) { bytes += static_cast(impl_->owner->size_bytes()); } return bytes; } -roaring_allowlist_view roaring_allowlist::view(std::size_t allowlist_id) const +roaring_allowlist_view roaring_allowlist::view() const noexcept { - RAFT_EXPECTS(allowlist_id < num_allowlists(), "Roaring allowlist_id is out of range."); - return roaring_allowlist_view{ - impl_->reference(allowlist_id), dataset_rows(), cardinality(allowlist_id)}; + return roaring_allowlist_view{impl_->reference(), dataset_rows(), cardinality()}; } -void roaring_allowlist::contains( - raft::resources const& res, - raft::device_matrix_view row_ids, - raft::device_matrix_view output) const +void roaring_allowlist::contains(raft::resources const& res, + raft::device_vector_view row_ids, + raft::device_vector_view output) const { contains_async(res, row_ids, output); raft::resource::sync_stream(res); @@ -2021,27 +237,19 @@ void roaring_allowlist::contains( void roaring_allowlist::contains_async( raft::resources const& res, - raft::device_matrix_view row_ids, - raft::device_matrix_view output) const -{ - RAFT_EXPECTS(row_ids.extent(0) == static_cast(num_allowlists()), - "Roaring membership matrix must have one row per allowlist."); - RAFT_EXPECTS(output.extent(0) == row_ids.extent(0) && output.extent(1) == row_ids.extent(1), - "Roaring membership output shape must match the input shape."); - auto const rows = static_cast(row_ids.extent(0)); - auto const columns = static_cast(row_ids.extent(1)); - if (columns == 0) { return; } - RAFT_EXPECTS(rows <= std::numeric_limits::max() / columns, - "Roaring membership matrix is too large."); - auto const size = rows * columns; - constexpr int block_size = 256; - contains_kernel<<>>( - impl_->references.data(), - impl_->empty_rows.data(), + raft::device_vector_view row_ids, + raft::device_vector_view output) const +{ + RAFT_EXPECTS(output.extent(0) == row_ids.extent(0), + "Roaring membership output size must match the input size."); + auto const size = static_cast(row_ids.extent(0)); + if (size == 0) { return; } + contains_kernel<<>>( + impl_->reference(), + empty(), static_cast(dataset_rows()), row_ids.data_handle(), output.data_handle(), - columns, size); RAFT_CUDA_TRY(cudaPeekAtLastError()); } diff --git a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu index e9f7534bea..047d246945 100644 --- a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu +++ b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu @@ -351,7 +351,7 @@ 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) { - // Equal tenant cardinalities let the same fixture exercise rectangular batched construction. + // 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) { @@ -384,111 +384,47 @@ TEST_P(CagraUdfFilterTest, TenantContextHonorsQuerySpecificMetadata) } } - // Reuse this existing query-specific UDF test as the exact reference for every single-partition - // algorithm, including max_queries=2 chunking above. - std::vector allowed_ids; - std::vector indptr{0}; - for (std::int64_t q = 0; q < n_queries; ++q) { - auto query_tenant = host_query_tenants[static_cast(q)]; + // 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)); } } - indptr.push_back(static_cast(allowed_ids.size())); + 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)); } - auto const tenant_width = indptr[1] - indptr[0]; - ASSERT_GT(tenant_width, 0); - for (std::int64_t query = 0; query < n_queries; ++query) { - ASSERT_EQ(indptr[static_cast(query + 1)] - indptr[static_cast(query)], - tenant_width); - } - auto tenant_allowlists = cuvs::core::roaring_allowlist::from_ids( - res, - n_rows, - raft::make_host_vector_view(allowed_ids.data(), - allowed_ids.size()), - raft::make_host_vector_view(indptr.data(), indptr.size())); std::vector tenant_views; tenant_views.reserve(n_queries); - for (std::int64_t query = 0; query < n_queries; ++query) { - tenant_views.push_back(tenant_allowlists.view(query)); + 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); - // Keep edge-selectivity coverage in this existing fixture: one filter mixes empty, full, - // fewer-than-k, sparse, and dense query allowlists. - std::vector> mixed_rows(static_cast(n_queries)); - for (std::uint32_t row = 0; row < n_rows; ++row) { - mixed_rows[1].push_back(row); - } - mixed_rows[2] = {7}; - mixed_rows[3] = {10, 11, 12}; - for (std::uint32_t row = 0; row < n_rows; row += 8) { - mixed_rows[4].push_back(row); - } - for (std::uint32_t row = n_rows - 32; row < n_rows; ++row) { - mixed_rows[5].push_back(row); - } - - std::vector mixed_ids; - std::vector mixed_indptr{0}; - for (auto const& row : mixed_rows) { - mixed_ids.insert(mixed_ids.end(), row.begin(), row.end()); - mixed_indptr.push_back(static_cast(mixed_ids.size())); - } - auto mixed_allowlists = cuvs::core::roaring_allowlist::from_ids( - res, - n_rows, - raft::make_host_vector_view(mixed_ids.data(), - mixed_ids.size()), - raft::make_host_vector_view(mixed_indptr.data(), - mixed_indptr.size()), - true); - std::vector mixed_views; - mixed_views.reserve(mixed_rows.size()); - for (std::size_t query = 0; query < mixed_rows.size(); ++query) { - mixed_views.push_back(mixed_allowlists.view(query)); - } - cuvs::neighbors::filtering::roaring_filter mixed_filter(res, mixed_views); - auto mixed_result = search(mixed_filter); - for (std::int64_t query = 0; query < n_queries; ++query) { - auto const& allowed = mixed_rows[static_cast(query)]; - for (std::int64_t rank = 0; rank < k; ++rank) { - auto row = mixed_result.neighbors[static_cast(query * k + rank)]; - auto valid_row = row < static_cast(n_rows); - if (query == 0) { EXPECT_FALSE(valid_row); } - if (query == 1) { EXPECT_TRUE(valid_row); } - if (valid_row) { EXPECT_NE(std::find(allowed.begin(), allowed.end(), row), allowed.end()); } - } - } - if (GetParam() == cagra::search_algo::SINGLE_CTA) { - std::array const one_empty{0, 0}; auto wrong_queries = cuvs::core::roaring_allowlist::from_ids( - res, - n_rows, - raft::make_host_vector_view(nullptr, 0), - raft::make_host_vector_view(one_empty.data(), - one_empty.size())); - std::array wrong_query_views{wrong_queries.view(0)}; + 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); - std::vector const all_empty(static_cast(n_queries) + 1, 0); auto wrong_columns = cuvs::core::roaring_allowlist::from_ids( - res, - n_rows + 1, - raft::make_host_vector_view(nullptr, 0), - raft::make_host_vector_view(all_empty.data(), - all_empty.size())); - std::vector wrong_column_views; - for (std::int64_t query = 0; query < n_queries; ++query) { - wrong_column_views.push_back(wrong_columns.view(query)); - } + 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); } diff --git a/cpp/tests/neighbors/roaring_allowlist.cu b/cpp/tests/neighbors/roaring_allowlist.cu index df5662d98b..633bfabb3d 100644 --- a/cpp/tests/neighbors/roaring_allowlist.cu +++ b/cpp/tests/neighbors/roaring_allowlist.cu @@ -1,13 +1,11 @@ /* * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * SPDX-License-Identifier: Apache-2.0 */ #include #include -#include - #include #include #include @@ -19,191 +17,45 @@ #include -#include #include -#include -#include #include #include #include #include -#include #include #include namespace cuvs::core { namespace { -// These hand-built fixtures exercise import of the portable bytes themselves. -// Keep their field order aligned with -// https://github.com/RoaringBitmap/RoaringFormatSpec. -constexpr std::uint32_t kCookieNoRun = 12346; -constexpr std::uint32_t kCookieRun = 12347; -constexpr std::uint32_t kSingleContainerPayloadByteOffset = 16; - -void append_u16(std::vector& out, std::uint16_t value) -{ - out.push_back(static_cast(value & 0xffu)); - out.push_back(static_cast((value >> 8) & 0xffu)); -} - -void append_u32(std::vector& out, std::uint32_t value) -{ - for (int i = 0; i < 4; ++i) { - out.push_back(static_cast((value >> (8 * i)) & 0xffu)); - } -} - -std::uint32_t read_u32(std::vector const& bytes, std::size_t offset = 0) -{ - std::uint32_t value{}; - for (int i = 0; i < 4; ++i) { - value |= static_cast(std::to_integer(bytes[offset + i])) - << (8 * i); - } - return value; -} - -// One no-run array container: cookie, N, {key, cardinality - 1}, payload -// offset, values. -std::vector array_row(std::vector const& values) -{ - std::vector out; - append_u32(out, kCookieNoRun); - append_u32(out, 1); - append_u16(out, 0); - append_u16(out, static_cast(values.size() - 1)); - append_u32(out, kSingleContainerPayloadByteOffset); - for (auto value : values) { - append_u16(out, value); - } - return out; -} - -// One run container: run cookie, one-byte run bitmap, descriptor, then the run -// payload. A run-cookie row with fewer than four containers has no -// container-offset table. -std::vector run_row(std::uint16_t start, std::uint16_t length_minus_one) -{ - std::vector out; - append_u32(out, kCookieRun); - out.push_back(std::byte{1}); - append_u16(out, 0); - append_u16(out, length_minus_one); - append_u16(out, 1); - append_u16(out, start); - append_u16(out, length_minus_one); - return out; -} - -// One no-run bitmap container containing the 5,000 even values below 10,000. -std::vector bitmap_row() -{ - constexpr std::uint16_t cardinality = 5000; - std::vector out; - append_u32(out, kCookieNoRun); - append_u32(out, 1); - append_u16(out, 0); - append_u16(out, cardinality - 1); - append_u32(out, kSingleContainerPayloadByteOffset); - out.resize(kSingleContainerPayloadByteOffset + 8192, std::byte{0}); - for (std::uint32_t value = 0; value < 10000; value += 2) { - out[kSingleContainerPayloadByteOffset + value / 8] |= static_cast(1u << (value % 8)); - } - return out; -} - -roaring_allowlist from_ragged_ids(raft::resources const& res, - std::size_t dataset_rows, - std::vector const& ids, - std::vector const& indptr, - bool pre_sorted = false) +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()), - raft::make_host_vector_view(indptr.data(), indptr.size()), pre_sorted); } -roaring_allowlist from_device_ragged_ids(raft::resources const& res, - std::size_t dataset_rows, - std::vector const& ids, - std::vector const& indptr, - bool pre_sorted = false) +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 device_indptr = raft::make_device_vector(res, indptr.size()); - auto const stream = raft::resource::get_cuda_stream(res); + 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); - raft::update_device(device_indptr.data_handle(), indptr.data(), indptr.size(), stream); return roaring_allowlist::from_ids( res, dataset_rows, raft::make_device_vector_view(device_ids.data_handle(), ids.size()), - raft::make_device_vector_view(device_indptr.data_handle(), - indptr.size()), pre_sorted); } -roaring_allowlist from_ids(raft::resources const& res, - std::size_t dataset_rows, - std::vector const& ids, - bool pre_sorted = false) -{ - return from_ragged_ids( - res, dataset_rows, ids, {0, static_cast(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) -{ - return from_device_ragged_ids( - res, dataset_rows, ids, {0, static_cast(ids.size())}, pre_sorted); -} - -roaring_allowlist import_row(raft::resources const& res, - std::size_t dataset_rows, - std::vector const& bytes) -{ - std::array const offsets{0, bytes.size()}; - return roaring_allowlist::from_serialized( - res, - dataset_rows, - raft::make_host_vector_view( - bytes.data(), static_cast(bytes.size())), - raft::make_host_vector_view(offsets.data(), offsets.size())); -} - -std::vector copy_serialized_bytes(raft::resources const& res, - roaring_allowlist const& allowlist) -{ - using ref_type = cuco::experimental::roaring_bitmap_ref; - static_assert(std::is_trivially_copyable_v); - auto const view = allowlist.view(0); - EXPECT_FALSE(view.empty()); - EXPECT_NE(view.device_reference(), nullptr); - - std::array raw_reference{}; - auto const stream = raft::resource::get_cuda_stream(res); - RAFT_CUDA_TRY(cudaMemcpyAsync(raw_reference.data(), - view.device_reference(), - raw_reference.size(), - cudaMemcpyDeviceToHost, - stream)); - raft::resource::sync_stream(res); - auto const reference = std::bit_cast(raw_reference); - std::vector bytes(reference.size_bytes()); - RAFT_CUDA_TRY( - cudaMemcpyAsync(bytes.data(), reference.data(), bytes.size(), cudaMemcpyDeviceToHost, stream)); - raft::resource::sync_stream(res); - return bytes; -} - void expect_membership(raft::resources const& res, roaring_allowlist const& allowlist, std::vector const& row_ids, @@ -215,20 +67,12 @@ void expect_membership(raft::resources const& res, 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); + auto rows_view = raft::make_device_vector_view( + rows.data_handle(), rows.size()); if (asynchronous) { - allowlist.contains_async( - res, - raft::make_device_matrix_view( - rows.data_handle(), 1, static_cast(row_ids.size())), - raft::make_device_matrix_view( - output.data_handle(), 1, static_cast(expected.size()))); + allowlist.contains_async(res, rows_view, output.view()); } else { - allowlist.contains( - res, - raft::make_device_matrix_view( - rows.data_handle(), 1, static_cast(row_ids.size())), - raft::make_device_matrix_view( - output.data_handle(), 1, static_cast(expected.size()))); + allowlist.contains(res, rows_view, output.view()); } std::vector actual(expected.size()); @@ -237,418 +81,79 @@ void expect_membership(raft::resources const& res, EXPECT_EQ(actual, expected); } -void expect_batch_membership(raft::resources const& res, - roaring_allowlist const& allowlists, - std::size_t columns, - std::vector const& row_ids, - std::vector const& expected) -{ - ASSERT_EQ(row_ids.size(), expected.size()); - ASSERT_EQ(row_ids.size(), allowlists.num_allowlists() * columns); - auto device_rows = raft::make_device_vector(res, row_ids.size()); - auto output = raft::make_device_vector(res, expected.size()); - auto const stream = raft::resource::get_cuda_stream(res); - raft::update_device(device_rows.data_handle(), row_ids.data(), row_ids.size(), stream); - allowlists.contains( - res, - raft::make_device_matrix_view( - device_rows.data_handle(), - static_cast(allowlists.num_allowlists()), - static_cast(columns)), - raft::make_device_matrix_view( - output.data_handle(), - static_cast(allowlists.num_allowlists()), - static_cast(columns))); - 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, BuildsArrayBitmapAndMultiContainerRowsFromIds) +TEST(RoaringAllowlist, BuildsArrayBitmapAndMultipleContainers) { raft::device_resources res; - auto array = from_ids(res, 200000, {65537, 7, 3, 131074, 5}); + 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(0), 5); - EXPECT_FALSE(array.empty(0)); + 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 consecutive; - for (std::uint32_t id = 100; id < 300; ++id) { - consecutive.push_back(id); - } - auto contiguous = from_ids(res, 1000, consecutive); - EXPECT_EQ(contiguous.cardinality(0), 200); - expect_membership(res, contiguous, {99, 100, 199, 299, 300}, {0, 1, 1, 1, 0}); - - std::vector sparse; + std::vector dense; for (std::uint32_t id = 0; id < 10000; id += 2) { - sparse.push_back(id); - } - auto bitmap = from_ids(res, 10000, sparse); - EXPECT_EQ(bitmap.cardinality(0), 5000); - expect_membership(res, bitmap, {0, 1, 8192, 9998, 9999}, {1, 0, 1, 1, 0}, true); -} - -TEST(RoaringAllowlist, BuildsRaggedRowsInOneGeneralBatch) -{ - raft::device_resources res; - constexpr std::size_t rows = 4; - constexpr std::size_t dataset_rows = (std::size_t{64} << 16) + 16; - std::vector ids; - std::vector indptr{0}; - - for (std::uint32_t i = 0; i < 64; ++i) { - ids.push_back(i); - } - indptr.push_back(static_cast(ids.size())); - indptr.push_back(static_cast(ids.size())); // empty row - for (std::uint32_t i = 0; i < 128; ++i) { - ids.push_back(1000 + 2 * i); - } - indptr.push_back(static_cast(ids.size())); - for (std::uint32_t i = 0; i < 64; ++i) { - ids.push_back((i << 16) + 7); - } - indptr.push_back(static_cast(ids.size())); - for (std::size_t row = 0; row < rows; ++row) { - std::reverse(ids.begin() + indptr[row], ids.begin() + indptr[row + 1]); + dense.push_back(id); } - - auto allowlists = from_ragged_ids(res, dataset_rows, ids, indptr); - auto device_allowlists = from_device_ragged_ids(res, dataset_rows, ids, indptr); - EXPECT_EQ(allowlists.num_allowlists(), rows); - EXPECT_EQ(allowlists.dataset_rows(), dataset_rows); - EXPECT_EQ(allowlists.total_cardinality(), ids.size()); - EXPECT_EQ(allowlists.cardinality(0), 64); - EXPECT_TRUE(allowlists.empty(1)); - EXPECT_EQ(allowlists.cardinality(2), 128); - EXPECT_EQ(allowlists.cardinality(3), 64); - EXPECT_EQ(allowlists.view(1).device_reference(), nullptr); - EXPECT_NE(allowlists.view(0).device_reference(), allowlists.view(2).device_reference()); - EXPECT_EQ(device_allowlists.total_cardinality(), ids.size()); - EXPECT_EQ(device_allowlists.size_bytes(), allowlists.size_bytes()); - EXPECT_EQ(read_u32(copy_serialized_bytes(res, allowlists)), kCookieNoRun); - - expect_batch_membership(res, - allowlists, - 4, - {0, - 63, - 64, - 1000, - 0, - 1000, - dataset_rows - 1, - dataset_rows, - 999, - 1000, - 1254, - 1255, - 7, - (std::uint32_t{63} << 16) + 7, - 63, - dataset_rows}, - {1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 0}); - expect_batch_membership(res, - device_allowlists, - 2, - {63, 64, 1, 1000, 1254, 1255, 7, (std::uint32_t{63} << 16) + 7}, - {1, 0, 0, 0, 1, 0, 1, 1}); - - std::array views{allowlists.view(0), allowlists.view(1), allowlists.view(2), allowlists.view(3)}; - cuvs::neighbors::filtering::roaring_filter filter(res, views); - EXPECT_EQ(filter.num_queries(), rows); - EXPECT_EQ(filter.cardinality(2), 128); + 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, BuildsGeneralMixedContainerRowsWithOneSegmentedSort) +TEST(RoaringAllowlist, SupportsPresortedHostAndDeviceInputs) { raft::device_resources res; - constexpr std::size_t rows = 3; - constexpr std::size_t width = 5000; - constexpr std::size_t dataset_rows = std::size_t{8} << 16; - std::vector const indptr{0, width, 2 * width, 3 * width}; - std::vector sorted; - sorted.reserve(rows * width); - for (std::uint32_t i = 0; i < width; ++i) { - sorted.push_back(100 + i); - } - for (std::uint32_t i = 0; i < width; ++i) { - sorted.push_back(2 * i); - } - for (std::uint32_t i = 0; i < width; ++i) { - auto const key = i % 8; - auto const value = (i / 8) * 2 + 1; - sorted.push_back((key << 16) + value); - } - for (std::size_t row = 0; row < rows; ++row) { - std::sort(sorted.begin() + indptr[row], sorted.begin() + indptr[row + 1]); - } - auto unsorted = sorted; - for (std::size_t row = 0; row < rows; ++row) { - std::reverse(unsorted.begin() + indptr[row], unsorted.begin() + indptr[row + 1]); - } + 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); - auto general = from_ragged_ids(res, dataset_rows, unsorted, indptr, false); - auto presorted = from_ragged_ids(res, dataset_rows, sorted, indptr, true); - EXPECT_EQ(general.num_allowlists(), rows); - EXPECT_EQ(general.total_cardinality(), rows * width); - EXPECT_EQ(general.size_bytes(), presorted.size_bytes()); - expect_batch_membership(res, - general, - 5, - {99, - 100, - 5099, - 5100, - 9998, - 0, - 1, - 8192, - 9998, - 9999, - 1, - 2, - (std::uint32_t{7} << 16) + 1249, - (std::uint32_t{7} << 16) + 1250, - dataset_rows}, - {0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0}); + 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}, true); + expect_membership(res, device, {0, 1, 7, 8, 65539}, {0, 1, 1, 0, 1}); } -TEST(RoaringAllowlist, ImportsRaggedPortableRowsIntoOnePackedOwner) +TEST(RoaringAllowlist, HandlesEmptyBoundsAndFullUint32Domain) { raft::device_resources res; - auto array = array_row({1, 3, 5}); - auto run = run_row(100, 99); - std::vector bytes; - bytes.insert(bytes.end(), array.begin(), array.end()); - auto const empty_offset = bytes.size(); - bytes.insert(bytes.end(), run.begin(), run.end()); - std::array const offsets{0, array.size(), empty_offset, bytes.size()}; - auto allowlists = roaring_allowlist::from_serialized( - res, - 1000, - raft::make_host_vector_view(bytes.data(), bytes.size()), - raft::make_host_vector_view(offsets.data(), offsets.size())); - EXPECT_EQ(allowlists.num_allowlists(), 3); - EXPECT_EQ(allowlists.cardinality(0), 3); - EXPECT_TRUE(allowlists.empty(1)); - EXPECT_EQ(allowlists.cardinality(2), 100); - EXPECT_EQ(allowlists.view(1).device_reference(), nullptr); - expect_batch_membership(res, - allowlists, - 4, - {1, 2, 3, 5, 1, 3, 100, 999, 99, 100, 199, 200}, - {1, 0, 1, 1, 0, 0, 0, 0, 0, 1, 1, 0}); -} -TEST(RoaringAllowlist, CucoConstructionPathsEmitByteIdenticalPortableRows) -{ - raft::device_resources res; - constexpr std::uint32_t second_key = std::uint32_t{1} << 16; - constexpr std::uint32_t third_key = std::uint32_t{2} << 16; - - std::vector ids; - for (std::uint32_t id = 100; id < 300; ++id) { - ids.push_back(id); // run - } - ids.insert(ids.end(), {second_key + 1, second_key + 3, second_key + 5}); // array - for (std::uint32_t id = 0; id < 10000; id += 2) { - ids.push_back(third_key + id); // bitmap - } - std::reverse(ids.begin(), ids.end()); - auto sorted = ids; - std::sort(sorted.begin(), sorted.end()); - auto const dataset_rows = static_cast(third_key) + 10000; - auto host_unsorted = from_ids(res, dataset_rows, ids); - auto host_pre_sorted = from_ids(res, dataset_rows, sorted, true); - auto device_unsorted = from_device_ids(res, dataset_rows, ids); - auto device_sorted = from_device_ids(res, dataset_rows, sorted, true); - - auto const expected = copy_serialized_bytes(res, host_unsorted); - EXPECT_EQ(copy_serialized_bytes(res, host_pre_sorted), expected); - EXPECT_EQ(copy_serialized_bytes(res, device_unsorted), expected); - EXPECT_EQ(copy_serialized_bytes(res, device_sorted), expected); -} - -TEST(RoaringAllowlist, BuildsFromUnsortedUniqueDeviceIdsWithoutModifyingInput) -{ - raft::device_resources res; - constexpr std::uint32_t second_key = std::uint32_t{1} << 16; - constexpr std::uint32_t third_key = std::uint32_t{2} << 16; - - std::vector ids; - for (std::uint32_t id = 100; id < 300; ++id) { - ids.push_back(id); // contiguous array container - } - ids.insert(ids.end(), {second_key + 1, second_key + 3, second_key + 5}); // array container - for (std::uint32_t id = 0; id < 10000; id += 2) { - ids.push_back(third_key + id); // bitmap container - } - std::reverse(ids.begin(), ids.end()); - auto const original = ids; - - auto device_ids = raft::make_device_vector(res, ids.size()); - auto device_indptr = raft::make_device_vector(res, 2); - std::array const indptr{0, static_cast(ids.size())}; - auto const stream = raft::resource::get_cuda_stream(res); - raft::update_device(device_ids.data_handle(), ids.data(), ids.size(), stream); - raft::update_device(device_indptr.data_handle(), indptr.data(), indptr.size(), stream); - auto allowlist = - roaring_allowlist::from_ids(res, - third_key + 10000, - raft::make_device_vector_view( - device_ids.data_handle(), ids.size()), - raft::make_device_vector_view( - device_indptr.data_handle(), 2)); - - std::vector unchanged(ids.size()); - raft::update_host(unchanged.data(), device_ids.data_handle(), unchanged.size(), stream); - raft::resource::sync_stream(res); - EXPECT_EQ(unchanged, original); - EXPECT_EQ(allowlist.cardinality(0), 5203); - expect_membership(res, - allowlist, - {99, - 100, - 299, - 300, - second_key + 1, - second_key + 2, - third_key, - third_key + 1, - third_key + 9998, - third_key + 9999}, - {0, 1, 1, 0, 1, 0, 1, 0, 1, 0}); -} - -TEST(RoaringAllowlist, PreSortedFastPathBuildsHostAndDeviceIds) -{ - raft::device_resources res; - std::vector const sorted_ids{1, 2, 3, 65537, 65539, 131072}; - auto host_allowlist = from_ids(res, 131073, sorted_ids, true); - auto device_allowlist = from_device_ids(res, 131073, sorted_ids, true); - EXPECT_EQ(host_allowlist.cardinality(0), sorted_ids.size()); - EXPECT_EQ(device_allowlist.cardinality(0), sorted_ids.size()); - EXPECT_EQ(host_allowlist.size_bytes(), device_allowlist.size_bytes()); - expect_membership(res, host_allowlist, {0, 1, 3, 4, 65537, 65538, 131072}, {0, 1, 1, 0, 1, 0, 1}); - expect_membership( - res, device_allowlist, {0, 1, 3, 4, 65537, 65538, 131072}, {0, 1, 1, 0, 1, 0, 1}); - - auto empty = from_device_ids(res, 10, {}, true); - EXPECT_TRUE(empty.empty(0)); - EXPECT_THROW(from_device_ids(res, 10, {0, 9, 10}, true), raft::logic_error); -} - -TEST(RoaringAllowlist, DeviceFactoryHandlesEmptyAndRejectsOutOfRangeIds) -{ - raft::device_resources res; - auto empty = from_device_ids(res, 10, {}); - EXPECT_TRUE(empty.empty(0)); - EXPECT_EQ(empty.size_bytes(), sizeof(void const*) + sizeof(std::uint8_t)); - expect_membership(res, empty, {0, 9, 10}, {0, 0, 0}); - - EXPECT_THROW(from_device_ids(res, 10, {0, 9, 10}), raft::logic_error); - EXPECT_THROW(from_device_ids(res, 10, {0, 65536, 131072}), raft::logic_error); - std::vector general_invalid(129); - for (std::uint32_t i = 0; i < 128; ++i) { - general_invalid[i] = i; - } - general_invalid.back() = std::numeric_limits::max(); - EXPECT_THROW(from_device_ids(res, 1000, general_invalid), raft::logic_error); -} - -TEST(RoaringAllowlist, SparseDeviceBuilderHandlesThresholdAndGeneralCrossover) -{ - raft::device_resources res; - constexpr auto second_key = std::uint32_t{1} << 16; - - std::vector sorted; - sorted.reserve(129); - for (std::uint32_t id = 100; id < 164; ++id) { - sorted.push_back(id); // one contiguous array container - } - for (std::uint32_t value = 1; value < 65; value += 2) { - sorted.push_back(second_key + value); // one array container - } - for (std::uint32_t key = 2; sorted.size() < 129; ++key) { - sorted.push_back((key << 16) + 7); // many one-value array containers - } - ASSERT_EQ(sorted.size(), 129); - - std::vector threshold_ids(sorted.begin(), sorted.begin() + 128); - auto reversed_threshold = threshold_ids; - std::reverse(reversed_threshold.begin(), reversed_threshold.end()); - auto reversed_general = sorted; - std::reverse(reversed_general.begin(), reversed_general.end()); - auto const dataset_rows = static_cast(sorted.back()) + 1; - - auto sparse_unsorted = from_device_ids(res, dataset_rows, reversed_threshold); - auto general_unsorted = from_device_ids(res, dataset_rows, reversed_general); - std::vector pre_sorted_ids(sorted.begin(), sorted.begin() + 64); - auto sparse_pre_sorted = from_device_ids(res, dataset_rows, pre_sorted_ids, true); - auto general_pre_sorted = from_device_ids(res, dataset_rows, sorted, true); + 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_EQ(sparse_unsorted.cardinality(0), 128); - EXPECT_EQ(general_unsorted.cardinality(0), 129); - EXPECT_EQ(sparse_pre_sorted.cardinality(0), 64); - EXPECT_EQ(general_pre_sorted.cardinality(0), 129); - EXPECT_EQ(general_unsorted.size_bytes(), general_pre_sorted.size_bytes()); + EXPECT_THROW(from_ids(res, 10, {10}), raft::logic_error); + EXPECT_THROW(from_ids(res, 0, {}), raft::logic_error); - auto const sparse_last = threshold_ids.back(); - auto const general_last = sorted.back(); - expect_membership(res, - sparse_unsorted, - {99, 100, 163, 164, second_key + 1, second_key + 2, sparse_last, general_last}, - {0, 1, 1, 0, 1, 0, 1, 0}); - expect_membership(res, - general_unsorted, - {99, 100, 163, 164, second_key + 1, second_key + 2, sparse_last, general_last}, - {0, 1, 1, 0, 1, 0, 1, 1}); - expect_membership(res, sparse_pre_sorted, {99, 100, 163, 164}, {0, 1, 1, 0}); - expect_membership( - res, general_pre_sorted, {99, 100, 163, 164, second_key + 1, general_last}, {0, 1, 1, 0, 1, 1}); -} - -TEST(RoaringAllowlist, SparseDeviceSortPreservesMaximumUint32Id) -{ - raft::device_resources res; - auto const max_id = std::numeric_limits::max(); - auto allowlist = - from_device_ids(res, static_cast(std::uint64_t{1} << 32), {max_id, 0}); - EXPECT_EQ(allowlist.cardinality(0), 2); - expect_membership(res, allowlist, {0, 1, max_id - 1, max_id}, {1, 0, 0, 1}); + 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(0); - auto next_view = allowlist.view(0); - EXPECT_TRUE(first_view.valid()); + auto first_view = allowlist.view(); + auto next_view = allowlist.view(); EXPECT_EQ(first_view.device_reference(), next_view.device_reference()); - EXPECT_NE(first_view.device_reference(), nullptr); auto moved = std::move(allowlist); - auto moved_view = moved.view(0); + 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}); - - auto empty = from_ids(res, 32, {}); - EXPECT_TRUE(empty.view(0).valid()); - EXPECT_EQ(empty.view(0).device_reference(), nullptr); - EXPECT_TRUE(empty.empty(0)); - EXPECT_EQ(empty.size_bytes(), sizeof(void const*) + sizeof(std::uint8_t)); - expect_membership(res, empty, {0, 31, 32}, {0, 0, 0}); } -TEST(RoaringAllowlist, StreamOrderedConstructionSupportsExplicitEventHandoff) +TEST(RoaringAllowlist, StreamOrderedConstructionSupportsEventHandoff) { rmm::cuda_stream build_stream; rmm::cuda_stream consume_stream; @@ -658,171 +163,46 @@ TEST(RoaringAllowlist, StreamOrderedConstructionSupportsExplicitEventHandoff) 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, raft::resource::get_cuda_stream(build_res))); - RAFT_CUDA_TRY(cudaStreamWaitEvent(raft::resource::get_cuda_stream(consume_res), ready)); - expect_membership(consume_res, allowlist, {99, 100, 200, 900, 901}, {0, 1, 1, 1, 0}, true); - - // Serialized import has the same lifetime rule. Keep the caller-owned host bytes alive until the - // event recorded after construction has completed, then consume the initialized reference on the - // second stream. - auto bytes = array_row({2, 4, 8}); - auto imported = import_row(build_res, 1000, bytes); - RAFT_CUDA_TRY(cudaEventRecord(ready, raft::resource::get_cuda_stream(build_res))); - RAFT_CUDA_TRY(cudaStreamWaitEvent(raft::resource::get_cuda_stream(consume_res), ready)); - expect_membership(consume_res, imported, {1, 2, 4, 7, 8}, {0, 1, 1, 0, 1}, true); - - // Releasing an owner immediately is safe: device_uvector schedules the packed allocation's - // deallocation after encoding and reference initialization on the construction stream. - { - auto temporary = from_ids(build_res, 1000, {1, 10, 100}); - EXPECT_NE(temporary.view(0).device_reference(), nullptr); - } - RAFT_CUDA_TRY(cudaEventRecord(ready, raft::resource::get_cuda_stream(build_res))); - RAFT_CUDA_TRY(cudaEventSynchronize(ready)); + 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(RoaringAllowlist, ImportsPortableArrayRunAndEmptyRows) -{ - raft::device_resources res; - - auto array_bytes = array_row({1, 3, 5}); - auto array = import_row(res, 1000, array_bytes); - EXPECT_EQ(array.cardinality(0), 3); - expect_membership(res, array, {1, 3, 4, 5}, {1, 1, 0, 1}); - - auto run_bytes = run_row(100, 99); - auto run = import_row(res, 1000, run_bytes); - EXPECT_EQ(run.cardinality(0), 100); - expect_membership(res, run, {99, 100, 199, 200}, {0, 1, 1, 0}, true); - - auto bitmap_bytes = bitmap_row(); - auto bitmap = import_row(res, 10000, bitmap_bytes); - EXPECT_EQ(bitmap.cardinality(0), 5000); - expect_membership(res, bitmap, {0, 1, 8192, 9998, 9999}, {1, 0, 1, 1, 0}, true); - - std::vector standard_empty; - append_u32(standard_empty, kCookieNoRun); - append_u32(standard_empty, 0); - auto empty = import_row(res, 1000, standard_empty); - EXPECT_TRUE(empty.empty(0)); - EXPECT_EQ(empty.view(0).device_reference(), nullptr); - - auto zero_length = import_row(res, 1000, {}); - EXPECT_TRUE(zero_length.empty(0)); -} - -TEST(RoaringAllowlist, RejectsSmallMalformedPortableInputs) -{ - raft::device_resources res; - - auto truncated = array_row({1, 3, 5}); - truncated.pop_back(); - EXPECT_THROW(import_row(res, 1000, truncated), raft::logic_error); - - auto bad_cookie = array_row({1}); - bad_cookie[0] = std::byte{0}; - EXPECT_THROW(import_row(res, 1000, bad_cookie), raft::logic_error); - - auto out_of_range = array_row({1000}); - EXPECT_THROW(import_row(res, 1000, out_of_range), raft::logic_error); - - auto valid = array_row({1, 3, 5}); - auto bytes = - raft::make_host_vector_view(valid.data(), valid.size()); - std::array const too_short{0}; - EXPECT_THROW(roaring_allowlist::from_serialized( - res, - 1000, - bytes, - raft::make_host_vector_view(too_short.data(), - too_short.size())), - raft::logic_error); - std::array const nonzero_start{1, valid.size()}; - EXPECT_THROW(roaring_allowlist::from_serialized( - res, - 1000, - bytes, - raft::make_host_vector_view( - nonzero_start.data(), nonzero_start.size())), - raft::logic_error); - std::array const decreasing{0, valid.size(), valid.size() - 1, valid.size()}; - EXPECT_THROW(roaring_allowlist::from_serialized( - res, - 1000, - bytes, - raft::make_host_vector_view(decreasing.data(), - decreasing.size())), - raft::logic_error); -} - -TEST(RoaringAllowlist, RejectsIdsOutsideLogicalShape) -{ - raft::device_resources res; - EXPECT_THROW(from_ids(res, 10, {10}), raft::logic_error); - EXPECT_THROW(from_ids(res, 0, {}), raft::logic_error); -} - -TEST(RoaringAllowlist, RejectsMalformedIndptr) -{ - raft::device_resources res; - std::vector const ids{1, 2}; - EXPECT_THROW(from_ragged_ids(res, 10, ids, {1, 2}), raft::logic_error); - EXPECT_THROW(from_ragged_ids(res, 10, ids, {0, 2, 1, 2}), raft::logic_error); - EXPECT_THROW(from_ragged_ids(res, 10, ids, {0, 1}), raft::logic_error); -} - -TEST(RoaringFilter, ReusesViewsAndUpdatesOneQueryOutsideSearch) +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(0), second.view(0), first.view(0)}; + 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); - EXPECT_EQ(filter.cardinality(2), 2); EXPECT_FLOAT_EQ(filter.filtering_rate(), 0.875f); - EXPECT_GT(filter.size_bytes(), 0); auto const* payload = filter.device_payload(); auto shared_copy = filter; - shared_copy.set_allowlist(res, 1, empty.view(0)); + 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); - shared_copy.set_allowlist(res, 1, second.view(0)); - EXPECT_EQ(filter.cardinality(1), 3); - EXPECT_EQ(filter.device_payload(), payload); -} - -TEST(RoaringFilter, RejectsSmallInvalidMappings) -{ - raft::device_resources res; EXPECT_THROW(cuvs::neighbors::filtering::roaring_filter( res, std::span{}), raft::logic_error); - - auto ten = from_ids(res, 10, {1}); - auto eleven = from_ids(res, 11, {1}); - std::array mismatched{ten.view(0), eleven.view(0)}; + 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); - - std::array invalid{cuvs::core::roaring_allowlist_view{}}; - EXPECT_THROW(cuvs::neighbors::filtering::roaring_filter(res, invalid), raft::logic_error); - - std::array one{ten.view(0)}; - cuvs::neighbors::filtering::roaring_filter filter(res, one); - EXPECT_THROW(filter.set_allowlist(res, 1, ten.view(0)), raft::logic_error); - EXPECT_THROW(filter.set_allowlist(res, 0, eleven.view(0)), 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 diff --git a/examples/cpp/src/cagra_roaring_filter_example.cu b/examples/cpp/src/cagra_roaring_filter_example.cu index 1a65d940c6..4835af001f 100644 --- a/examples/cpp/src/cagra_roaring_filter_example.cu +++ b/examples/cpp/src/cagra_roaring_filter_example.cu @@ -21,54 +21,10 @@ 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; -constexpr std::uint32_t portable_cookie_no_run = 12346; -constexpr std::uint32_t single_array_payload_byte_offset = 16; - -void append_u16(std::vector& bytes, std::uint16_t value) -{ - bytes.push_back(static_cast(value & 0xffu)); - bytes.push_back(static_cast((value >> 8) & 0xffu)); -} - -void append_u32(std::vector& bytes, std::uint32_t value) -{ - for (int i = 0; i < 4; ++i) { - bytes.push_back(static_cast((value >> (8 * i)) & 0xffu)); - } -} - -// This helper intentionally writes the simplest nonempty portable Roaring row: -// -// bytes 0..3: uint32 no-run cookie (12346) -// bytes 4..7: uint32 container count (1) -// bytes 8..11: { uint16 key, uint16 cardinality_minus_one } -// bytes 12..15: uint32 payload offset, measured from byte 0 of this row -// bytes 16..: sorted uint16 array values -// -// Every integer is little-endian. This example's IDs are below 2^16, so they all use container key -// zero; a general encoder must split IDs by their high 16 bits and may need bitmap or run -// containers. Prefer `from_ids` when starting with IDs. `from_serialized` is intended primarily -// for interoperating with systems that already emit the standard format: -// -// https://github.com/RoaringBitmap/RoaringFormatSpec -void append_portable_array_row(std::vector& bytes, - std::vector const& ids, - std::int64_t first, - std::int64_t last) -{ - append_u32(bytes, portable_cookie_no_run); - append_u32(bytes, 1); // one container - append_u16(bytes, 0); // high 16-bit container key - append_u16(bytes, static_cast(last - first - 1)); - append_u32(bytes, single_array_payload_byte_offset); - for (auto i = first; i < last; ++i) { - append_u16(bytes, static_cast(ids[static_cast(i)])); - } -} +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 @@ -91,51 +47,29 @@ int main() 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); - index.update_device_dataset_same_layout(res, padded); - // Ragged construction: query q accepts rows whose row id modulo n_queries equals q. - std::vector ids; - std::vector indptr{0}; + // 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); } - indptr.push_back(static_cast(ids.size())); - } - // One factory call consumes the ragged rows directly and retains every variable-length stream - // in one owner. - auto id_allowlists = cuvs::core::roaring_allowlist::from_ids( - res, - n_rows, - raft::make_host_vector_view(ids.data(), ids.size()), - raft::make_host_vector_view(indptr.data(), indptr.size()), - true); - std::vector id_views; - for (std::int64_t query = 0; query < n_queries; ++query) { - id_views.push_back(id_allowlists.view(query)); + owners.push_back(cuvs::core::roaring_allowlist::from_ids( + res, + n_rows, + raft::make_host_vector_view(ids.data(), ids.size()), + true)); } - cuvs::neighbors::filtering::roaring_filter ids_filter(res, id_views); - // A serialized system can keep ragged portable rows in one host buffer plus outer offsets. The - // whole packed buffer becomes one owner even though encoded row sizes may differ. - std::vector bytes; - std::vector byte_offsets{0}; - for (std::int64_t query = 0; query < n_queries; ++query) { - append_portable_array_row(bytes, ids, indptr[query], indptr[query + 1]); - byte_offsets.push_back(bytes.size()); - } - - auto serialized_allowlists = cuvs::core::roaring_allowlist::from_serialized( - res, - n_rows, - raft::make_host_vector_view(bytes.data(), bytes.size()), - raft::make_host_vector_view(byte_offsets.data(), - byte_offsets.size())); - std::vector serialized_views; - for (std::int64_t query = 0; query < n_queries; ++query) { - serialized_views.push_back(serialized_allowlists.view(query)); + std::vector views; + views.reserve(owners.size()); + for (auto const& owner : owners) { + views.push_back(owner.view()); } - cuvs::neighbors::filtering::roaring_filter filter(res, serialized_views); + 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); @@ -167,8 +101,7 @@ int main() } } - std::cout << "Built " << id_allowlists.num_allowlists() << " ID allowlists and imported " - << serialized_allowlists.num_allowlists() - << " portable allowlists; CAGRA reused one prepared filter payload.\n"; - return ids_filter.num_queries() == n_queries ? 0 : 1; + std::cout << "CAGRA reused " << owners.size() + << " independently owned Roaring allowlists through one prepared filter payload.\\n"; + return 0; } From 2637cf4cd0227a1a9fbc91ce00502c81c9e786f1 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Wed, 2 Sep 2026 20:45:34 +0000 Subject: [PATCH 4/5] reduce surface area --- cpp/include/cuvs/core/roaring_allowlist.hpp | 16 +-- cpp/src/core/roaring_allowlist.cu | 134 ++++---------------- cpp/src/neighbors/roaring_filter.cu | 25 +--- cpp/tests/neighbors/roaring_allowlist.cu | 40 ++++-- 4 files changed, 63 insertions(+), 152 deletions(-) diff --git a/cpp/include/cuvs/core/roaring_allowlist.hpp b/cpp/include/cuvs/core/roaring_allowlist.hpp index bcf706989d..c90423eb2b 100644 --- a/cpp/include/cuvs/core/roaring_allowlist.hpp +++ b/cpp/include/cuvs/core/roaring_allowlist.hpp @@ -88,7 +88,8 @@ class CUVS_EXPORT roaring_allowlist { * @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. Empty input is valid and rejects every candidate. + * device overload. IDs must be smaller than dataset_rows; this precondition is not checked. + * Empty input is valid and rejects every candidate. */ static roaring_allowlist from_ids(raft::resources const& res, std::size_t dataset_rows, @@ -98,7 +99,8 @@ class CUVS_EXPORT roaring_allowlist { /** * @brief Build one allowlist from device IDs. * - * The input must remain valid until the construction stream reaches the enqueued work. + * 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. */ static roaring_allowlist from_ids(raft::resources const& res, @@ -123,16 +125,6 @@ class CUVS_EXPORT roaring_allowlist { /** @brief Return a zero-copy view. */ [[nodiscard]] roaring_allowlist_view view() const noexcept; - /** @brief Test row IDs and synchronize the resource stream. */ - void contains(raft::resources const& res, - raft::device_vector_view row_ids, - raft::device_vector_view output) const; - - /** @brief Stream-ordered asynchronous version of @ref contains. */ - void contains_async(raft::resources const& res, - raft::device_vector_view row_ids, - raft::device_vector_view output) const; - private: explicit roaring_allowlist(std::unique_ptr impl) noexcept; diff --git a/cpp/src/core/roaring_allowlist.cu b/cpp/src/core/roaring_allowlist.cu index 33a3d5a6f9..16631b56a0 100644 --- a/cpp/src/core/roaring_allowlist.cu +++ b/cpp/src/core/roaring_allowlist.cu @@ -9,10 +9,10 @@ #include +#include #include #include #include -#include #include #include @@ -20,11 +20,10 @@ #include #include -#include #include #include -#include #include +#include #include #include @@ -42,8 +41,8 @@ namespace { * * 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 materializes that reference in - * a stable device allocation during construction. Views and CAGRA filters copy only its pointer, so + * 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 @@ -54,14 +53,6 @@ using ref_type = cuco::experimental::roaring_bitmap_ref; using cuco_bitmap_type = cuco::experimental::roaring_bitmap; -constexpr int kBlockSize = 256; - -int grid_size_for(std::size_t count) -{ - auto const blocks = (count + kBlockSize - 1) / kBlockSize; - return static_cast(std::min(blocks, 65535)); -} - void validate_dataset_rows(std::size_t dataset_rows) { constexpr std::uint64_t kKeyDomain = std::uint64_t{1} << 32; @@ -70,100 +61,58 @@ void validate_dataset_rows(std::size_t dataset_rows) "dataset_rows exceeds the uint32_t Roaring key domain."); } -__global__ void validate_input_kernel(std::uint32_t const* ids, - std::size_t size, - std::uint64_t dataset_rows, - std::uint32_t* invalid) -{ - auto const i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (i < size && static_cast(ids[i]) >= dataset_rows) { atomicExch(invalid, 1u); } -} - -__global__ void store_ref_kernel(ref_type ref, ref_type* output) -{ - if (threadIdx.x == 0) { ::new (static_cast(output)) ref_type{ref}; } -} - -__global__ void contains_kernel(ref_type const* reference, - bool empty, - std::uint64_t dataset_rows, - 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) { return; } - auto const row = row_ids[i]; - output[i] = !empty && static_cast(row) < dataset_rows && reference->contains(row); -} - struct device_build_result { std::unique_ptr owner; - rmm::device_uvector reference; std::size_t cardinality{}; }; device_build_result build_from_device_ids( raft::resources const& res, - std::size_t dataset_rows, 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 {nullptr, rmm::device_uvector(0, stream), 0}; } - - rmm::device_uvector invalid(1, stream); - RAFT_CUDA_TRY(cudaMemsetAsync(invalid.data(), 0, sizeof(std::uint32_t), stream)); - validate_input_kernel<<>>( - ids.data_handle(), size, dataset_rows, invalid.data()); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - - std::uint32_t host_invalid{}; - RAFT_CUDA_TRY(cudaMemcpyAsync( - &host_invalid, invalid.data(), sizeof(host_invalid), cudaMemcpyDeviceToHost, stream)); + 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( + 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( + : cuco_bitmap_type::from_indices( ids.data_handle(), ids.data_handle() + size, allocator, cuco_stream); - - // The exact-size readback in the cuco factory orders the preceding validation copy. - RAFT_EXPECTS(host_invalid == 0, "Roaring allowlist ID must be smaller than dataset_rows."); - auto owner = std::make_unique(std::move(bitmap)); auto const cardinality = static_cast(owner->size()); - rmm::device_uvector reference(sizeof(ref_type), stream); - store_ref_kernel<<<1, 1, 0, stream>>>(owner->ref(), - reinterpret_cast(reference.data())); - RAFT_CUDA_TRY(cudaPeekAtLastError()); - return {std::move(owner), std::move(reference), cardinality}; + return {std::move(owner), cardinality}; } } // namespace struct roaring_allowlist::impl { std::unique_ptr owner; - rmm::device_uvector reference_storage; + std::optional host_reference; + rmm::device_uvector device_reference; std::size_t cardinality_{}; std::size_t dataset_rows_{}; - impl(std::size_t dataset_rows, device_build_result&& built) + impl(raft::resources const& res, std::size_t dataset_rows, device_build_result&& built) : owner(std::move(built.owner)), - reference_storage(std::move(built.reference)), + 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_destructible_v); + 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 cardinality_ == 0 ? nullptr - : reinterpret_cast(reference_storage.data()); + return device_reference.size() == 0 ? nullptr : device_reference.data(); } }; @@ -181,17 +130,11 @@ roaring_allowlist roaring_allowlist::from_ids( 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_CUDA_TRY(cudaMemcpyAsync(device_ids.data(), - ids.data_handle(), - size * sizeof(key_type), - cudaMemcpyHostToDevice, - 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, dataset_rows, device_ids_view, pre_sorted); - return roaring_allowlist{std::make_unique(dataset_rows, std::move(built))}; + 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( @@ -201,8 +144,8 @@ roaring_allowlist roaring_allowlist::from_ids( bool pre_sorted) { validate_dataset_rows(dataset_rows); - auto built = build_from_device_ids(res, dataset_rows, ids, pre_sorted); - return roaring_allowlist{std::make_unique(dataset_rows, std::move(built))}; + 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; @@ -217,7 +160,7 @@ bool roaring_allowlist::empty() const noexcept { return cardinality() == 0; } std::size_t roaring_allowlist::size_bytes() const noexcept { - auto bytes = impl_->reference_storage.size() * sizeof(cuda::std::byte); + auto bytes = impl_->device_reference.size() * sizeof(ref_type); if (impl_->owner) { bytes += static_cast(impl_->owner->size_bytes()); } return bytes; } @@ -227,31 +170,4 @@ roaring_allowlist_view roaring_allowlist::view() const noexcept return roaring_allowlist_view{impl_->reference(), dataset_rows(), cardinality()}; } -void roaring_allowlist::contains(raft::resources const& res, - raft::device_vector_view row_ids, - raft::device_vector_view output) const -{ - contains_async(res, row_ids, output); - raft::resource::sync_stream(res); -} - -void roaring_allowlist::contains_async( - raft::resources const& res, - raft::device_vector_view row_ids, - raft::device_vector_view output) const -{ - RAFT_EXPECTS(output.extent(0) == row_ids.extent(0), - "Roaring membership output size must match the input size."); - auto const size = static_cast(row_ids.extent(0)); - if (size == 0) { return; } - contains_kernel<<>>( - impl_->reference(), - empty(), - static_cast(dataset_rows()), - row_ids.data_handle(), - output.data_handle(), - size); - RAFT_CUDA_TRY(cudaPeekAtLastError()); -} - } // namespace cuvs::core diff --git a/cpp/src/neighbors/roaring_filter.cu b/cpp/src/neighbors/roaring_filter.cu index 148186443f..9079593fea 100644 --- a/cpp/src/neighbors/roaring_filter.cu +++ b/cpp/src/neighbors/roaring_filter.cu @@ -8,15 +8,13 @@ #include #include +#include #include #include #include -#include #include -#include - #include #include #include @@ -89,22 +87,13 @@ struct roaring_filter::impl { host_empty.push_back(allowlist.empty() ? 1 : 0); } - RAFT_CUDA_TRY(cudaMemcpyAsync(refs.data(), - host_refs.data(), - host_refs.size() * sizeof(ref_type const*), - cudaMemcpyHostToDevice, - stream)); - RAFT_CUDA_TRY(cudaMemcpyAsync(empty_rows.data(), - host_empty.data(), - host_empty.size() * sizeof(std::uint8_t), - cudaMemcpyHostToDevice, - stream)); + 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_CUDA_TRY(cudaMemcpyAsync( - payload.data(), &host_payload, sizeof(host_payload), cudaMemcpyHostToDevice, stream)); + 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. @@ -172,10 +161,8 @@ void roaring_filter::set_allowlist(raft::resources const& res, 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_CUDA_TRY(cudaMemcpyAsync( - impl_->refs.data() + query_id, &ref, sizeof(ref), cudaMemcpyHostToDevice, stream)); - RAFT_CUDA_TRY(cudaMemcpyAsync( - impl_->empty_rows.data() + query_id, &empty_, sizeof(empty_), cudaMemcpyHostToDevice, stream)); + 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; diff --git a/cpp/tests/neighbors/roaring_allowlist.cu b/cpp/tests/neighbors/roaring_allowlist.cu index 633bfabb3d..970b03efd4 100644 --- a/cpp/tests/neighbors/roaring_allowlist.cu +++ b/cpp/tests/neighbors/roaring_allowlist.cu @@ -6,12 +6,15 @@ #include #include +#include + #include #include #include #include #include #include +#include #include @@ -56,24 +59,38 @@ roaring_allowlist from_device_ids(raft::resources const& res, 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, - bool asynchronous = false) + 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); - auto rows_view = raft::make_device_vector_view( - rows.data_handle(), rows.size()); - if (asynchronous) { - allowlist.contains_async(res, rows_view, output.view()); - } else { - allowlist.contains(res, rows_view, output.view()); - } + 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); @@ -115,11 +132,11 @@ TEST(RoaringAllowlist, SupportsPresortedHostAndDeviceInputs) 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}, true); + 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, HandlesEmptyBoundsAndFullUint32Domain) +TEST(RoaringAllowlist, HandlesEmptyAndFullUint32Domain) { raft::device_resources res; @@ -130,7 +147,6 @@ TEST(RoaringAllowlist, HandlesEmptyBoundsAndFullUint32Domain) EXPECT_EQ(empty.view().device_reference(), nullptr); expect_membership(res, empty, {0, 31, 32}, {0, 0, 0}); - EXPECT_THROW(from_ids(res, 10, {10}), raft::logic_error); EXPECT_THROW(from_ids(res, 0, {}), raft::logic_error); auto maximum = From b7fcf7880d9763efd8ee3ab81a25e34798d94f5c Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 10 Sep 2026 19:12:40 +0000 Subject: [PATCH 5/5] address review --- cpp/cmake/patches/cuco_override.json | 2 +- cpp/include/cuvs/core/roaring_allowlist.hpp | 5 +++- cpp/include/cuvs/neighbors/common.hpp | 11 +++++++-- .../cuvs/neighbors/dynamic_batching.hpp | 3 ++- cpp/include/cuvs/neighbors/tiered_index.hpp | 3 ++- cpp/src/neighbors/detail/dynamic_batching.cuh | 17 ++++++++++++-- cpp/src/neighbors/detail/tiered_index.cuh | 4 ++++ cpp/src/neighbors/roaring_filter.cu | 2 ++ .../neighbors/ann_cagra/test_filter_udf.cu | 23 +++++++++++++++++++ cpp/tests/neighbors/roaring_allowlist.cu | 1 + 10 files changed, 63 insertions(+), 8 deletions(-) diff --git a/cpp/cmake/patches/cuco_override.json b/cpp/cmake/patches/cuco_override.json index 55204afb14..004b881cb2 100644 --- a/cpp/cmake/patches/cuco_override.json +++ b/cpp/cmake/patches/cuco_override.json @@ -3,7 +3,7 @@ "cuco" : { "version": "0.0.1", "git_url": "https://github.com/NVIDIA/cuCollections.git", - "git_tag": "9d7c9307395c3b8795d93ad65d0751c98471dde6" + "git_tag": "2027216da1ba687374222f54e0e196dd1ccaedc7" } } } diff --git a/cpp/include/cuvs/core/roaring_allowlist.hpp b/cpp/include/cuvs/core/roaring_allowlist.hpp index c90423eb2b..af27678252 100644 --- a/cpp/include/cuvs/core/roaring_allowlist.hpp +++ b/cpp/include/cuvs/core/roaring_allowlist.hpp @@ -89,7 +89,8 @@ class CUVS_EXPORT roaring_allowlist { * * 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. + * 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, @@ -102,6 +103,8 @@ class CUVS_EXPORT roaring_allowlist { * 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, diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 83ec04e345..04155f37d2 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -1529,6 +1529,10 @@ struct bloom_filter : public base_filter { * 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 */ @@ -1556,9 +1560,12 @@ struct roaring_filter : public base_filter { [[nodiscard]] bool empty(std::size_t query_id) const; /** - * @brief Maximum rejected fraction among all query allowlists. + * @brief Conservative maximum rejected fraction among all query allowlists. * - * CAGRA uses this precomputed value when `search_params::filtering_rate` is unset. + * 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; 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/neighbors/detail/dynamic_batching.cuh b/cpp/src/neighbors/detail/dynamic_batching.cuh index 9839a7fe19..323b57222a 100644 --- a/cpp/src/neighbors/detail/dynamic_batching.cuh +++ b/cpp/src/neighbors/detail/dynamic_batching.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -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/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 index 9079593fea..bfd56a6948 100644 --- a/cpp/src/neighbors/roaring_filter.cu +++ b/cpp/src/neighbors/roaring_filter.cu @@ -49,6 +49,8 @@ std::size_t validate_views(std::span a 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()); diff --git a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu index 047d246945..e510b11709 100644 --- a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu +++ b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu @@ -414,6 +414,29 @@ TEST_P(CagraUdfFilterTest, TenantContextHonorsQuerySpecificMetadata) 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)); diff --git a/cpp/tests/neighbors/roaring_allowlist.cu b/cpp/tests/neighbors/roaring_allowlist.cu index 970b03efd4..5e07416926 100644 --- a/cpp/tests/neighbors/roaring_allowlist.cu +++ b/cpp/tests/neighbors/roaring_allowlist.cu @@ -202,6 +202,7 @@ TEST(RoaringFilter, ReusesViewsUpdatesMappingsAndRejectsInvalidInputs) 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();