diff --git a/README.md b/README.md index fa5885281..f9ec0ecf1 100644 --- a/README.md +++ b/README.md @@ -267,4 +267,4 @@ We plan to add many GPU-accelerated, concurrent data structures to `cuCollection `cuco::experimental::roaring_bitmap` implements a Roaring bitmap following the [Roaring bitmap format specification](https://github.com/RoaringBitmap/RoaringFormatSpec). #### Examples: -- [Host-bulk APIs](https://github.com/NVIDIA/cuCollections/blob/dev/examples/roaring_bitmap/host_bulk_example.cu) (see [live example in godbolt](https://godbolt.org/clientstate/eJy9WAtPGzkQ_itzi1RtIMmGlEcbHteUlCq6HlSUe0hQrZxdJ7GyWW9tL5BD_Pcb2_uEpdDHXZDIrj3-5uFvxuPcOpJKyXgsncHFrcNCZ7DZdiISz1Iyo87ACdKQOG1H8lQE-t1bv4xhHT59HP3dOWYRPeLJSrDZXJ3TGzWA4hXcoAX9Xn-7g_922nDy53g0HsLR6dnH07Ph-fj0BF7A8Ph4_GE8PH_3qQvDKAKzUoKgkoorGnZLVR9YQGNJO-OQxopNGRUDGCYkmNNOv9vTct5lfBmvsTiI0pDCfpAG3BOcCBbP_AlTS5J0g3R--EAmVSxiauUpQZiS3XmSHN5HCoknVegF-C-k08NHJ1msmifVKqG-VVATUHORSuWF9Ard869ooLjozptEIj5jAYmaJ9OYXVEhSVSFqMpNcaPkSiq6rC2fSiUoqY8x3jCIQxjG2pDVZPR465YTb7QamCOAP0mjhU9vyDKJKIbdTk8Eo1MY0SWyDYOhqIRUIsuAT0HNKdR3Cy4djXLpQMT5Ik0MMAw_jqWhhYEcx7iQScg0wTVFYRKCuubwst9BILBgEkgcAo8p7GxVhsFNuFBkgkunXCyJasFU8KW2xuBfnFmT3hrpYyPyKaHBZ3euVCIHnjdjap5OugFfejXZ_K1c00JaJ1wyjNpKW2MUIH-DBTDrv3Y3dxb9RHmViliauYALgRHXqZFGSFM4IUsardraZQykMkJTHkX8GrWC2fCBUdGBC-vsNZrKUyXSWHYnLDaTro1S63sc8iYRn3jbm7u7JHztaSNCoojXqKz10JT_z44HRuSbnjFtZ6u0w9LjJ9qxs-U1qmsVJH4Tc0XhvMpjQb-kDLfabv2SLDBHEoVVGjqjoz-OTv3R6V8nH06HI__sdHg2PnnvYwk9Hw3PhwdYVxWHCQVJVZEoWBsx95MIcw6LBhajGNkDv9HVOT5jDk84jywXXbR6MLD5jqTDRH2R5YqvOeUnRM3R9FtEBZKiphmNqc5lf0FXEg7g4rPbgs4h2NI0GNRq236uEgwAaOYbJfQmEXhiYL3UytECJn2JNvpX-ZI2VGZTrLQv-746bOVAAJ4HR1i40MMvKcUUM_ZgVtvUKshwi6y8yxlR5kkO8dPod_ZuOPr9XXcZrumhjh7L1RgXsoA0OWVM38ultQtugxgsMNq9Pfzah82e_ujnjQPzUokLGLhuksq5PyG4w4tWgX1XU4LABrRE28cT3D5vbCy-hvkSefwk7m4F99XzcJswbVl8hGBN8XQN6ITOWOy22lYFjUO3lYPfAY0k_REy7mx9Axkb64Gh4n_BRCxAJRcfqLaF7ylaGveepqUWs7S8MdubRpHdbXx_nb9_x44_pWtY17XZ-wFd90R6N_0MrZRtFNmuiTxqMbNJy-4ZynTi9jXEV8zt3bzKFsAGsB_MDLOlz82M25IjRLHAJxIbdOXq_lmrwfNEt-X-lKBwUePb2MFlz7BEg_ShVBQvjE9uxaVTupJ5cHuX69df-kV_Y1KcoiqTKfrd5mbWxJpR9_5R1c6EuBwMkOpErKwuTHX3Fy3VxdzmCIou516aFQEVAvb30YVjgmKhPlS1HA7o4ft69JhZh5GIMtszX0xQzNBd4cZ7qoy9INk_tDhKDZoeQY4YsLJtt89m8oGPrTI8eIfiga46pldeYqeNtad-BhbRy1lSXh5M556wOKZhA2kmK4Xn9iSdTqlwC2sqys8odt7GK9xYnmnXcybOuEehKyhOUZFgaPyASLUfzIlYP3RzWwS59hNuZMy8a9V1dR3DHUKK1hRn0EHEJXUrlmSlN7tPFI5nnX0RAUtfrPZUsCXyl0RoQe0OUvYr9XHX7vBzzd6rbLztlqrHgkswoSV2y1GoUwSPIEUY7gGGsbJvrYImWaNVa7wavCdg91AzV-I3tXcse3-oEqBeJnQreFgaYSuE4Z2OvuZyRdVbvLF0shtLxSMyw7XZncRar-XvX8atCvloDWqXVuTTVS_zm5OOnlZkNYtHA6jXmT4XF_ilzEERBz3Op-4DpVU7MsOwPlS6A_OThFrd3llOZnlfU4MTJvt1K7QkGtX00GtsGuKV2LT0eSevu3h_ND4rKly1GQdNKT9kAu1uXLVXuCnTIKBSQv1zYNv85j6pAN_AMtd8kcuLdY7-4rsAvxXNHlZ1tOZLnJNTxFZxtNxW8UzDIKvftqJhmEiUzIkZyUy4V8grG5pL_Ao9GMAmTq7pw7FUVhwZzXe5xv3KLmvIAsNGrBwMy4COUf3ed-mUbUH2uXSecRls3V_44KDKvLP-xCGbGqo6bQc7zgQrpSh_GXTiqyDY7G-nmzhtDcNJp4OAB8HGxuYudIgI5gdy6e_2oNPB0qrwn9LtQdiJyHJifkuM2KSCGQRBhIP6DEI8HNBcWzh37Xweq3RtHuuVc_fZ_P0LZpjzEQ==)) \ No newline at end of file +- [Host-bulk APIs](https://github.com/NVIDIA/cuCollections/blob/dev/examples/roaring_bitmap/host_bulk_example.cu) (see [live example in godbolt](https://godbolt.org/clientstate/eJy9WAtP20gQ_itzRqocSOJAebQpcE1JqaLrQUW5h1Qqa2NvklUcr7u7BnIo__1md_0MptDHXZCIvTs7z29mZ3LnSCol47F0-p_uHBY6_e22E5F4mpIpdfpOkIbEaTuSpyLQ797mVQyb8PHD8O_OKYvoCU-Wgk1n6pLeqj4Ur-AGLdjp7ex18N9-G87-HA1HAzg5v_hwfjG4HJ2fwTMYnJ6O3o8Gl28_dmEQRWBOShBUUnFNw24p6j0LaCxpZxTSWLEJo6IPg4QEM9rZ6fY0nXcVX8UbLA6iNKRwGKQB9wQngsVTf8zUgiTdIJ0d36NJFYuYWnpKEKZkd5Ykx-ucQuJJFXoB_gvp5PjBTRar5k2ZkLh5Ry0T6lvRNQI1E6lUXkiv0XD_mgaKi-6siSTiUxaQqHkzjdk1FZJEVRZVugmGUC6loova8YlUgpL6GuMNi7iEDq4tWUlGjrdp0fJai4EZMvDHaTT36S1ZJBHFgNjtsWB0AkO6QByiMxSVkErEH_AJqBmFehzhytFcrhyIOJ-niWEMgw8jaQBjWI5iPMgkZJLghiIxCUHdcHi-00FGYJlJIHEIPKawv1tZBjfhQpExHp1wsSCqBRPBF1obw__ThVXpjaE-NSQfExp8dmdKJbLveVOmZum4G_CFV6PN38ozLQR8wiVDry21NkYAIjuYA7P2a3NzY9FOpFepiKXZC7gQ6HGdNGmEAIYzsqDRsq1NRkcqQzThUcRvUCqYgPeNiA58ssbeoKo8VSKNZXfMYrPpWi-1vscgbxzxsbe3fXBAwpeeViIkiniNwlr3Vfn_9LinRB70DGn7u6UeFh4_UY_9Xa9RXKsA8euYKwqXVRwL-iVlGGob-gWZY44kCus3dIYnf5yc-8Pzv87enw-G_sX54GJ09s7H4no5HFwOjrDiKg5jCpKqIlGwamLuJxHmHBYNLEYxogd-o8tLfMYcHnMeWSy6qHW_b_MdQYeJ-izLFV9jyk-ImqHqd8gVSIqSpjSmOpf9OV1KOIJPn90WdI7BlqZ-v1bbDnORYBiARr4RQm8TgXcJ1kstHDVg0peoo3-dH2lDZTfFGvx8x1fHrZwRgOfBCRYutPBLSjHFjD6Y1Ta1CjDcISpXOSLKPMlZ_DT4XbwdDH9_212EG3qpo9dyMcaEzCFNRhnVX-XU2gS3gQzm6O3eK_w6hO2e_ujnrSPzUvELGHbdJJUzf0wwwvNWwXtVE4KMDdOS2yHe7fZ5a2v-NZ7PEceP8j2o8H3xNL5NPG1ZfABgTf50DdMxnbLYbbWtCBqHbitnvgIaSfojYNzf_QYwNtYDA8X_AolYgEos3hNtC99jsDTmPQ5LTWZheWvCm0aRjTa-v8zfvyPij8ka1GVt935A1hpJ73Yn41bSNpLs1Uge1JjZpGVrijKduDuaxVfU7d2-yA7AFrAfzAwT0qdmxl2JEaJY4BOJrbtydWetxeB9oht2f0KQuKjxbezgsmdYoEL6UiqKF_on1-LKKU3JLLhb5fL1l37R35gU5yjKZIp-t7mZNbFm1V2_qtoZEZf9PkKdiKWVhanu_qKpupjbHJmiybmV5kRAhYDDQzThlCBZqC9VTYcLenldjl4z59ATUaZ7ZotxillaFWa8o8roC5L9Q4ur1HDTK4gRw6xs2-2z2bxnY6t0D05XPNBVx_TKC-y0sfbU78DCezlKyuHBdO4Ji2MaNoBmvFR4b4_TyYQKt9CmIvyCYudtrMLA8ky63jN-xhiFrqC4RUWCrvEDItVhMCNi89jNdRHkxk-4oTH7rhXX1XUMI4QQrQnOWAcRl9StaJKV3myeKAzPOvvCAxa-WO2pYAvEL4lQg9oMUvYr9XXXRrjiHT37rXvL3ibHd0-0rp05t2ui3FpVDHqXNVnV28QlWAckNtlRqDMLZSnCMHTo_Uq4WwW6sv6s1q81OI2ADb0GvMRvakczO3ZUcVOvLrqDPC6VsIXFGtK2KVAR9QYHnU426FQsIlM8m40yVntNvz7dWxHywdLVLrXIt6tW5gOX9p4WZCWLBx2oz5n2GA_4Jc1R4Qe9zifuPaFVPTLFsKxUmgrzG4da3q0slLNyURODG6Zo6A5qQTRX03pvsEmIk7SZBPIBQDf__nB0URTGag8PGmJ-yATq3XjqVWGmTIOASgn1z5GdDprbq4L5FlbH5vkvr_E592ffxfBbudk7rs6tefZzcojY4o-a2-KfSehnZd-mNrqJRMmMmJVMhbX6XwloTvEr9KAP27i5oe_UUlhx0zSPgI3xymY8RIFBI1YShmVA-6g-Ll45ZTeRfa6cJ8yQrfWD9-63zDprTxyyiYGq03awUU2wwIryp0Ynvg6C7Z29dBu3rWK46XSQ4VGwtbV9AB0igtmRXPgHPeh0sCIr_Kd0VxF2IrIYmx8nIzau8AyCIMJFfXUhP1zQWJs7q3a-j8W9to_1yll9Nn__Ar14EEI=)) \ No newline at end of file diff --git a/benchmarks/roaring_bitmap/contains_bench.cu b/benchmarks/roaring_bitmap/contains_bench.cu index 77c9cef31..8ca870c1a 100644 --- a/benchmarks/roaring_bitmap/contains_bench.cu +++ b/benchmarks/roaring_bitmap/contains_bench.cu @@ -13,6 +13,7 @@ #include #include +#include #include #include @@ -40,7 +41,8 @@ void roaring_bitmap_contains(nvbench::state& state, nvbench::type_list) file.read(reinterpret_cast(thrust::raw_pointer_cast(buffer.data())), file_size); file.close(); - cuco::experimental::roaring_bitmap roaring_bitmap(thrust::raw_pointer_cast(buffer.data())); + cuco::experimental::roaring_bitmap roaring_bitmap( + cuda::std::span{thrust::raw_pointer_cast(buffer.data()), buffer.size()}); thrust::device_vector items(num_items); diff --git a/examples/roaring_bitmap/host_bulk_example.cu b/examples/roaring_bitmap/host_bulk_example.cu index d64986c61..2ae3d7236 100644 --- a/examples/roaring_bitmap/host_bulk_example.cu +++ b/examples/roaring_bitmap/host_bulk_example.cu @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -95,7 +96,7 @@ bool check(std::string const& bitmap_file_path) // Create roaring bitmap from the file cuco::experimental::roaring_bitmap roaring_bitmap( - thrust::raw_pointer_cast(buffer.data())); + cuda::std::span{thrust::raw_pointer_cast(buffer.data()), buffer.size()}); // Generate query keys (all should be contained in the bitmap) auto keys = generate_keys(); diff --git a/include/cuco/detail/roaring_bitmap/roaring_bitmap.inl b/include/cuco/detail/roaring_bitmap/roaring_bitmap.inl index 61326fdd3..da15f7e83 100644 --- a/include/cuco/detail/roaring_bitmap/roaring_bitmap.inl +++ b/include/cuco/detail/roaring_bitmap/roaring_bitmap.inl @@ -6,10 +6,19 @@ #pragma once #include +#include #include namespace cuco::experimental { +template +roaring_bitmap::roaring_bitmap(cuda::std::span bitmap, + Allocator const& alloc, + cuda::stream_ref stream) + : storage_{bitmap, alloc, stream} +{ +} + template roaring_bitmap::roaring_bitmap(cuda::std::byte const* bitmap, Allocator const& alloc, diff --git a/include/cuco/detail/roaring_bitmap/roaring_bitmap_storage.cuh b/include/cuco/detail/roaring_bitmap/roaring_bitmap_storage.cuh index 0d8e465a3..88760a8b7 100644 --- a/include/cuco/detail/roaring_bitmap/roaring_bitmap_storage.cuh +++ b/include/cuco/detail/roaring_bitmap/roaring_bitmap_storage.cuh @@ -8,10 +8,12 @@ #include #include #include +#include #include #include #include +#include #include #include @@ -245,6 +247,27 @@ class roaring_bitmap_storage { ~roaring_bitmap_storage() = default; + /** + * @brief Constructs storage by validating and copying bitmap data to device memory + * + * @param bitmap Serialized bitmap bytes in host memory + * @param alloc Allocator for device memory allocation + * @param stream CUDA stream for memory operations + */ + roaring_bitmap_storage(cuda::std::span bitmap, + Allocator const& alloc, + cuda::stream_ref stream) + : allocator_{alloc}, + metadata_{bitmap}, + data_{allocator_.allocate(metadata_.size_bytes, stream), + cuco::detail::custom_deleter{ + metadata_.size_bytes, allocator_, stream}}, + ref_{data_.get(), metadata_} + { + CUCO_CUDA_TRY(cuco::detail::memcpy_async( + data_.get(), bitmap.data(), metadata_.size_bytes, cudaMemcpyHostToDevice, stream)); + } + /** * @brief Constructs storage by copying bitmap data to device memory * @@ -262,8 +285,8 @@ class roaring_bitmap_storage { metadata_.size_bytes, allocator_, stream}}, ref_{data_.get(), metadata_} { - CUCO_CUDA_TRY(cudaMemcpyAsync( - data_.get(), bitmap, metadata_.size_bytes, cudaMemcpyHostToDevice, stream.get())); + CUCO_CUDA_TRY(cuco::detail::memcpy_async( + data_.get(), bitmap, metadata_.size_bytes, cudaMemcpyHostToDevice, stream)); } /** @@ -335,6 +358,48 @@ class roaring_bitmap_storage { ~roaring_bitmap_storage() = default; + /** + * @brief Constructs storage by validating and copying bitmap data to device memory + * + * @param bitmap Serialized bitmap bytes in host memory + * @param alloc Allocator for device memory allocation + * @param stream CUDA stream for memory operations + */ + roaring_bitmap_storage(cuda::std::span bitmap, + Allocator const& alloc, + cuda::stream_ref stream) + : allocator_{alloc}, + bucket_allocator_{alloc}, + bucket_metadata_{}, + buckets_h_{}, + metadata_{ + [bitmap](std::vector& bucket_metadata) { + return typename ref_type::metadata_type{bitmap, bucket_metadata}; + }(bucket_metadata_)}, + data_{allocator_.allocate(metadata_.size_bytes, stream), + cuco::detail::custom_deleter{ + metadata_.size_bytes, allocator_, stream}}, + buckets_{bucket_allocator_.allocate(metadata_.num_buckets, stream), + cuco::detail::custom_deleter{ + metadata_.num_buckets, bucket_allocator_, stream}}, + ref_{data_.get(), metadata_, buckets_.get()} + { + assert(metadata_.valid); + buckets_h_.reserve(bucket_metadata_.size()); + for (auto const& meta : bucket_metadata_) { + buckets_h_.emplace_back(meta.key, + bucket_ref_type{data_.get() + meta.byte_offset, meta.metadata}); + } + CUCO_CUDA_TRY(cuco::detail::memcpy_async( + data_.get(), bitmap.data(), metadata_.size_bytes, cudaMemcpyHostToDevice, stream)); + CUCO_CUDA_TRY(cuco::detail::memcpy_async( + buckets_.get(), + buckets_h_.data(), + metadata_.num_buckets * sizeof(cuda::std::pair), + cudaMemcpyHostToDevice, + stream)); + } + /** * @brief Constructs storage by copying bitmap data to device memory * @@ -367,14 +432,14 @@ class roaring_bitmap_storage { buckets_h_.emplace_back(meta.key, bucket_ref_type{data_.get() + meta.byte_offset, meta.metadata}); } - CUCO_CUDA_TRY(cudaMemcpyAsync( - data_.get(), bitmap, metadata_.size_bytes, cudaMemcpyHostToDevice, stream.get())); - CUCO_CUDA_TRY(cudaMemcpyAsync( + CUCO_CUDA_TRY(cuco::detail::memcpy_async( + data_.get(), bitmap, metadata_.size_bytes, cudaMemcpyHostToDevice, stream)); + CUCO_CUDA_TRY(cuco::detail::memcpy_async( buckets_.get(), buckets_h_.data(), metadata_.num_buckets * sizeof(cuda::std::pair), cudaMemcpyHostToDevice, - stream.get())); + stream)); } /** diff --git a/include/cuco/detail/roaring_bitmap/util.cuh b/include/cuco/detail/roaring_bitmap/util.cuh index 87841f664..e2a2a69c9 100644 --- a/include/cuco/detail/roaring_bitmap/util.cuh +++ b/include/cuco/detail/roaring_bitmap/util.cuh @@ -11,7 +11,9 @@ #include #include #include +#include #include +#include #include #include @@ -40,6 +42,108 @@ __host__ __device__ __forceinline__ bool check_bit(cuda::std::byte const* bitmap (cuda::std::uint8_t(1) << (index % 8)); } +/** + * @brief Non-owning view of serialized bitmap data with optional bounds information + * + * Pointer-backed views preserve the unchecked behavior of the legacy API, while span-backed views + * validate accesses against the serialized data size. + */ +class serialized_bitmap_view { + public: + /** + * @brief Constructs an unbounded view from a pointer + * + * @param data Pointer to the beginning of the serialized bitmap + */ + __host__ __device__ explicit serialized_bitmap_view(cuda::std::byte const* data) + : data_{data}, size_{0}, bounded_{false} + { + } + + /** + * @brief Constructs a bounded view from a span + * + * @param bitmap Serialized bitmap bytes + */ + __host__ __device__ explicit serialized_bitmap_view(cuda::std::span bitmap) + : data_{bitmap.data()}, size_{bitmap.size()}, bounded_{true} + { + } + + /** + * @brief Returns a pointer to the serialized bitmap data + * + * @return Pointer to the beginning of the serialized bitmap + */ + [[nodiscard]] __host__ __device__ cuda::std::byte const* data() const noexcept { return data_; } + + /** + * @brief Returns the size of a bounded view + * + * @return Serialized bitmap size in bytes, or zero for an unbounded view + */ + [[nodiscard]] __host__ __device__ cuda::std::size_t size() const noexcept { return size_; } + + /** + * @brief Indicates whether the view has bounds information + * + * @return true if the view was constructed from a span, otherwise false + */ + [[nodiscard]] __host__ __device__ bool is_bounded() const noexcept { return bounded_; } + + /** + * @brief Checks whether a byte range is contained in the view + * + * Unbounded views contain every range. + * + * @param offset Start of the range in bytes + * @param size Size of the range in bytes + * @return true if the range is contained in the view, otherwise false + */ + [[nodiscard]] __host__ __device__ bool contains(cuda::std::size_t offset, + cuda::std::size_t size) const noexcept + { + return not bounded_ or (offset <= size_ and size <= size_ - offset); + } + + /** + * @brief Returns a view beginning at the specified byte offset + * + * @param offset Offset from the beginning of the serialized bitmap + * @return View of the remaining serialized bitmap data + */ + [[nodiscard]] __host__ __device__ serialized_bitmap_view subview(cuda::std::size_t offset) const + { + if (not bounded_) { return serialized_bitmap_view{data_ + offset}; } + if (offset > size_) { + return serialized_bitmap_view{cuda::std::span{data_, 0}}; + } + return serialized_bitmap_view{ + cuda::std::span{data_ + offset, size_ - offset}}; + } + + /** + * @brief Loads a value if its byte range is contained in the view + * + * @tparam T Type of value to load + * @param offset Offset of the value in bytes + * @param value Reference that receives the loaded value + * @return true if the value was loaded, otherwise false + */ + template + __host__ __device__ bool try_load(cuda::std::size_t offset, T& value) const + { + if (not contains(offset, sizeof(T))) { return false; } + value = misaligned_load(data_ + offset); + return true; + } + + private: + cuda::std::byte const* data_; + cuda::std::size_t size_; + bool bounded_; +}; + template struct roaring_bitmap_metadata { static_assert(cuco::dependent_false, "T must be either uint32_t or uint64_t"); @@ -81,25 +185,102 @@ struct roaring_bitmap_metadata { /// Whether container offsets are stored in the serialized data bool offsets_in_serialized_data = true; + /** + * @brief Constructs metadata from a bounded serialized bitmap + * + * @param bitmap Serialized bitmap bytes + */ + __host__ roaring_bitmap_metadata(cuda::std::span bitmap) + : roaring_bitmap_metadata{serialized_bitmap_view{bitmap}} + { + } + /** * @brief Constructs metadata from a serialized bitmap * * @param bitmap Pointer to the beginning of the serialized bitmap */ __host__ __device__ roaring_bitmap_metadata(cuda::std::byte const* bitmap) + : roaring_bitmap_metadata{serialized_bitmap_view{bitmap}} + { + } + + /** + * @brief Constructs metadata from an internal serialized bitmap view + * + * @param bitmap Serialized bitmap view + */ + __host__ __device__ explicit roaring_bitmap_metadata(serialized_bitmap_view bitmap) + { + parse(bitmap); + } + + private: + template + __host__ __device__ bool load(serialized_bitmap_view bitmap, cuda::std::size_t offset, T& value) + { + if (bitmap.try_load(offset, value)) { return true; } + valid = false; + NV_IF_TARGET( + NV_IS_HOST, + CUCO_FAIL("Invalid bitmap format: serialized data is truncated");) // TODO device error + // handling + return false; + } + + __host__ __device__ bool expect_range(serialized_bitmap_view bitmap, + cuda::std::size_t offset, + cuda::std::size_t size) + { + if (bitmap.contains(offset, size)) { return true; } + valid = false; + NV_IF_TARGET( + NV_IS_HOST, + CUCO_FAIL("Invalid bitmap format: serialized data is truncated");) // TODO device error + // handling + return false; + } + + __host__ __device__ bool get_container_size(serialized_bitmap_view bitmap, + cuda::std::size_t container_offset, + cuda::std::int32_t index, + cuda::std::size_t& size) + { + bool const is_run_container = + has_run and check_bit(bitmap.data() + run_container_bitmap, index); + if (is_run_container) { + cuda::std::uint16_t num_runs; + if (not load(bitmap, container_offset, num_runs)) { return false; } + size = sizeof(cuda::std::uint16_t) + + static_cast(num_runs) * 2 * sizeof(cuda::std::uint16_t); + return true; + } + + auto const card_offset = + static_cast(key_cards) + + static_cast(index * 2 + 1) * sizeof(cuda::std::uint16_t); + cuda::std::uint16_t stored_card; + if (not load(bitmap, card_offset, stored_card)) { return false; } + auto const card = 1u + stored_card; + size = card <= max_array_container_card + ? static_cast(card) * sizeof(cuda::std::uint16_t) + : static_cast(bitset_container_bytes); + return true; + } + + __host__ __device__ void parse(serialized_bitmap_view bitmap) { constexpr cuda::std::uint32_t serial_cookie_no_runcontainer = 12346; constexpr cuda::std::uint32_t serial_cookie = 12347; - // constexpr cuda::std::uint32_t frozen_cookie = 13766; // not implemented - constexpr cuda::std::int32_t max_containers = 1 << 16; - constexpr cuda::std::uint32_t cookie_mask = 0xFFFF; - constexpr cuda::std::uint32_t cookie_shift = 16; - - cuda::std::byte const* buf = bitmap; + constexpr cuda::std::uint32_t max_containers = 1 << 16; + constexpr cuda::std::uint32_t cookie_mask = 0xFFFF; + constexpr cuda::std::uint32_t cookie_shift = 16; + cuda::std::size_t offset = 0; cuda::std::uint32_t cookie; - cuda::std::memcpy(&cookie, buf, sizeof(cuda::std::uint32_t)); - buf += sizeof(cuda::std::uint32_t); + if (not load(bitmap, offset, cookie)) { return; } + offset += sizeof(cuda::std::uint32_t); + if ((cookie & cookie_mask) != serial_cookie && cookie != serial_cookie_no_runcontainer) { valid = false; NV_IF_TARGET( @@ -110,15 +291,14 @@ struct roaring_bitmap_metadata { return; } - if ((cookie & cookie_mask) == serial_cookie) - // upper 16 bits of cookie are the number of containers - 1 - num_containers = (cookie >> cookie_shift) + 1; - else { - // following 4 bytes are the number of containers - cuda::std::memcpy(&num_containers, buf, sizeof(cuda::std::uint32_t)); - buf += sizeof(cuda::std::uint32_t); + cuda::std::uint32_t container_count; + if ((cookie & cookie_mask) == serial_cookie) { + container_count = (cookie >> cookie_shift) + 1; + } else { + if (not load(bitmap, offset, container_count)) { return; } + offset += sizeof(cuda::std::uint32_t); } - if (num_containers < 0 or num_containers > max_containers) { + if (container_count > max_containers) { valid = false; NV_IF_TARGET( NV_IS_HOST, @@ -126,98 +306,116 @@ struct roaring_bitmap_metadata { "Invalid bitmap format: num_containers out of range");) // TODO device error handling return; } + num_containers = static_cast(container_count); has_run = (cookie & cookie_mask) == serial_cookie; if (has_run) { - cuda::std::size_t s = (num_containers + 7) / 8; // ceil bytes to store run container bitmap - run_container_bitmap = cuda::std::distance(bitmap, buf); - buf += s; + auto const run_container_bitmap_size = + static_cast((num_containers + 7) / 8); + if (not expect_range(bitmap, offset, run_container_bitmap_size)) { return; } + run_container_bitmap = static_cast(offset); + offset += run_container_bitmap_size; } - key_cards = cuda::std::distance(bitmap, buf); - // if the current address is aligned to 2 bytes, then all containers are aligned to at least 2 - // bytes - bool const aligned_16 = (reinterpret_cast(bitmap + key_cards) % - sizeof(cuda::std::uint16_t)) == 0; - buf += num_containers * 2 * sizeof(cuda::std::uint16_t); + key_cards = static_cast(offset); + auto const key_cards_size = + static_cast(num_containers) * 2 * sizeof(cuda::std::uint16_t); + if (not expect_range(bitmap, offset, key_cards_size)) { return; } + offset += key_cards_size; if ((!has_run) || (num_containers >= no_offset_threshold)) { - // Container offsets are stored in the serialized data offsets_in_serialized_data = true; - container_offsets = cuda::std::distance(bitmap, buf); - buf += num_containers * sizeof(cuda::std::uint32_t); + container_offsets = static_cast(offset); + auto const container_offsets_size = + static_cast(num_containers) * sizeof(cuda::std::uint32_t); + if (not expect_range(bitmap, offset, container_offsets_size)) { return; } + offset += container_offsets_size; } else { - // Container offsets are NOT stored in the serialized data - // We need to compute them by walking through the containers offsets_in_serialized_data = false; container_offsets = 0; + } - cuda::std::byte const* container_ptr = buf; - for (cuda::std::int32_t i = 0; i < num_containers; ++i) { - // Store the computed offset for this container - computed_offsets[i] = - static_cast(cuda::std::distance(bitmap, container_ptr)); - - // Get cardinality for this container - cuda::std::byte const* card_ptr = - bitmap + key_cards + (i * 2 + 1) * sizeof(cuda::std::uint16_t); - cuda::std::uint32_t card_i = 1u + misaligned_load(card_ptr); - - // Check if this is a run container - bool is_run_container = check_bit(bitmap + run_container_bitmap, i); - - // Compute container size and advance pointer - if (is_run_container) { - // Run container: first uint16_t is num_runs, followed by num_runs (start, length) pairs - cuda::std::uint16_t num_runs = misaligned_load(container_ptr); - container_ptr += sizeof(cuda::std::uint16_t) + num_runs * 2 * sizeof(cuda::std::uint16_t); - } else if (card_i <= max_array_container_card) { - // Array container - container_ptr += card_i * sizeof(cuda::std::uint16_t); - } else { - // Bitset container (fixed size) - container_ptr += bitset_container_bytes; - } - } - // buf now points past all containers - buf = container_ptr; + if (num_containers == 0) { + size_bytes = offset; + valid = true; + return; } - cuda::std::uint32_t card = 0; - for (cuda::std::int32_t i = 0; i < num_containers; i++) { - cuda::std::byte const* card_ptr = - bitmap + key_cards + (i * 2 + 1) * sizeof(cuda::std::uint16_t); - if (aligned_16) { - card = 1u + aligned_load(card_ptr); - } else { - card = 1u + misaligned_load(card_ptr); + for (cuda::std::int32_t i = 0; i < num_containers; ++i) { + auto const card_offset = + static_cast(key_cards) + + static_cast(i * 2 + 1) * sizeof(cuda::std::uint16_t); + cuda::std::uint16_t stored_card; + if (not load(bitmap, card_offset, stored_card)) { return; } + auto const card = 1u + stored_card; + if (card > cuda::std::numeric_limits::max() - num_keys) { + valid = false; + NV_IF_TARGET( + NV_IS_HOST, + CUCO_FAIL("Invalid bitmap format: cardinality overflow");) // TODO device error handling + return; } num_keys += card; } - // find end of roaring bitmap (re-use card from last container) - cuda::std::byte const* end; if (offsets_in_serialized_data) { - end = - bitmap + misaligned_load( - bitmap + container_offsets + (num_containers - 1) * sizeof(cuda::std::uint32_t)); - } else { - end = bitmap + computed_offsets[num_containers - 1]; - } - - if (has_run and check_bit(bitmap + run_container_bitmap, num_containers - 1)) { - cuda::std::uint16_t const num_runs = misaligned_load(end); - end += sizeof(cuda::std::uint16_t) + num_runs * 2 * sizeof(cuda::std::uint16_t); - } else { - if (card <= max_array_container_card) { - end += card * sizeof(cuda::std::uint16_t); + if (not bitmap.is_bounded()) { + auto const last_container = num_containers - 1; + auto const offset_offset = + static_cast(container_offsets) + + static_cast(last_container) * sizeof(cuda::std::uint32_t); + cuda::std::uint32_t stored_offset; + if (not load(bitmap, offset_offset, stored_offset)) { return; } + auto const container_offset = static_cast(stored_offset); + cuda::std::size_t size; + if (not get_container_size(bitmap, container_offset, last_container, size)) { return; } + size_bytes = container_offset + size; } else { - end += bitset_container_bytes; // fixed size bitset container + auto const containers_start = offset; + auto previous_end = containers_start; + for (cuda::std::int32_t i = 0; i < num_containers; ++i) { + auto const offset_offset = + static_cast(container_offsets) + + static_cast(i) * sizeof(cuda::std::uint32_t); + cuda::std::uint32_t stored_offset; + if (not load(bitmap, offset_offset, stored_offset)) { return; } + auto const container_offset = static_cast(stored_offset); + if (container_offset < containers_start or container_offset < previous_end) { + valid = false; + NV_IF_TARGET( + NV_IS_HOST, + CUCO_FAIL("Invalid bitmap format: container offsets are invalid");) // TODO device + // error handling + return; + } + cuda::std::size_t size; + if (not get_container_size(bitmap, container_offset, i, size)) { return; } + if (not expect_range(bitmap, container_offset, size)) { return; } + previous_end = container_offset + size; + } + size_bytes = previous_end; } + } else { + for (cuda::std::int32_t i = 0; i < num_containers; ++i) { + if (offset > cuda::std::numeric_limits::max()) { + valid = false; + NV_IF_TARGET( + NV_IS_HOST, + CUCO_FAIL( + "Invalid bitmap format: container offset is out of range");) // TODO device error + // handling + return; + } + computed_offsets[i] = static_cast(offset); + cuda::std::size_t size; + if (not get_container_size(bitmap, offset, i, size)) { return; } + if (not expect_range(bitmap, offset, size)) { return; } + offset += size; + } + size_bytes = offset; } - size_bytes = static_cast(cuda::std::distance(bitmap, end)); - valid = true; + valid = true; } }; @@ -266,6 +464,18 @@ struct roaring_bitmap_metadata { } }; + /** + * @brief Constructs metadata from a bounded serialized 64-bit bitmap + * + * @param bitmap Serialized bitmap bytes + * @param bucket_metadata Vector to store metadata for each bucket + */ + __host__ roaring_bitmap_metadata(cuda::std::span bitmap, + std::vector& bucket_metadata) + { + parse(serialized_bitmap_view{bitmap}, bucket_metadata); + } + /** * @brief Constructs metadata from a serialized 64-bit bitmap with bucket metadata * @@ -275,52 +485,114 @@ struct roaring_bitmap_metadata { __host__ roaring_bitmap_metadata(cuda::std::byte const* bitmap, std::vector& bucket_metadata) { - cuda::std::size_t byte_offset = 0; - cuda::std::byte const* bitmap_ptr = bitmap; - cuda::std::memcpy(&num_buckets, bitmap_ptr, sizeof(cuda::std::uint64_t)); - byte_offset += sizeof(cuda::std::uint64_t); // skip num_buckets + parse(serialized_bitmap_view{bitmap}, bucket_metadata); + } + + /** + * @brief Constructs metadata from a serialized 64-bit bitmap + * + * @param bitmap Pointer to the beginning of the serialized bitmap + */ + __host__ __device__ roaring_bitmap_metadata(cuda::std::byte const* bitmap) + { + parse(serialized_bitmap_view{bitmap}); + } + + private: + template + __host__ __device__ bool load(serialized_bitmap_view bitmap, cuda::std::size_t offset, T& value) + { + if (bitmap.try_load(offset, value)) { return true; } + valid = false; + NV_IF_TARGET( + NV_IS_HOST, + CUCO_FAIL("Invalid bitmap format: serialized data is truncated");) // TODO device error + // handling + return false; + } + + __host__ void parse(serialized_bitmap_view bitmap, std::vector& bucket_metadata) + { + cuda::std::size_t byte_offset = 0; + cuda::std::uint64_t serialized_num_buckets; + if (not load(bitmap, byte_offset, serialized_num_buckets)) { return; } + byte_offset += sizeof(cuda::std::uint64_t); + + CUCO_EXPECTS(serialized_num_buckets <= cuda::std::numeric_limits::max(), + "Invalid bitmap format: num_buckets out of range"); + num_buckets = static_cast(serialized_num_buckets); + + constexpr cuda::std::size_t minimum_bucket_prefix_size = + sizeof(cuda::std::uint32_t) + sizeof(cuda::std::uint32_t); + if (bitmap.is_bounded()) { + CUCO_EXPECTS(num_buckets <= (bitmap.size() - byte_offset) / minimum_bucket_prefix_size, + "Invalid bitmap format: num_buckets exceeds the serialized data size"); + } bucket_metadata.clear(); - bucket_metadata.reserve(num_buckets); + if (not bitmap.is_bounded()) { bucket_metadata.reserve(num_buckets); } for (cuda::std::size_t i = 0; i < num_buckets; ++i) { cuda::std::uint32_t bucket_key; - cuda::std::memcpy(&bucket_key, bitmap_ptr + byte_offset, sizeof(cuda::std::uint32_t)); - byte_offset += sizeof(cuda::std::uint32_t); // skip bucket key - roaring_bitmap_metadata bucket_meta{bitmap_ptr + byte_offset}; - if (!bucket_meta.valid) { - valid = false; - return; - } + if (not load(bitmap, byte_offset, bucket_key)) { return; } + byte_offset += sizeof(cuda::std::uint32_t); + + roaring_bitmap_metadata bucket_meta{bitmap.subview(byte_offset)}; + CUCO_EXPECTS(bucket_meta.valid and bucket_meta.size_bytes > 0, + "Invalid bitmap format: bucket metadata is invalid"); + CUCO_EXPECTS(not bitmap.is_bounded() or bitmap.contains(byte_offset, bucket_meta.size_bytes), + "Invalid bitmap format: serialized data is truncated"); + CUCO_EXPECTS( + bucket_meta.num_keys <= cuda::std::numeric_limits::max() - num_keys, + "Invalid bitmap format: cardinality overflow"); + CUCO_EXPECTS( + bucket_meta.size_bytes <= cuda::std::numeric_limits::max() - byte_offset, + "Invalid bitmap format: bitmap size overflow"); + bucket_metadata.emplace_back(byte_offset, bucket_key, bucket_meta); num_keys += bucket_meta.num_keys; - byte_offset += bucket_meta.size_bytes; // skip bucket + byte_offset += bucket_meta.size_bytes; } size_bytes = byte_offset; valid = true; } - /** - * @brief Constructs metadata from a serialized 64-bit bitmap - * - * @param bitmap Pointer to the beginning of the serialized bitmap - */ - __host__ __device__ roaring_bitmap_metadata(cuda::std::byte const* bitmap) + __host__ __device__ void parse(serialized_bitmap_view bitmap) { - cuda::std::size_t byte_offset = 0; - cuda::std::byte const* bitmap_ptr = bitmap; - cuda::std::memcpy(&num_buckets, bitmap_ptr, sizeof(cuda::std::uint64_t)); - byte_offset += sizeof(cuda::std::uint64_t); // skip num_buckets + cuda::std::size_t byte_offset = 0; + cuda::std::uint64_t serialized_num_buckets; + if (not load(bitmap, byte_offset, serialized_num_buckets)) { return; } + if (serialized_num_buckets > cuda::std::numeric_limits::max()) { + valid = false; + NV_IF_TARGET(NV_IS_HOST, + CUCO_FAIL("Invalid bitmap format: num_buckets out of range");) // TODO device + // error handling + return; + } + num_buckets = static_cast(serialized_num_buckets); + byte_offset += sizeof(cuda::std::uint64_t); for (cuda::std::size_t i = 0; i < num_buckets; ++i) { - byte_offset += sizeof(cuda::std::uint32_t); // skip bucket key - roaring_bitmap_metadata bucket_meta{bitmap_ptr + byte_offset}; - if (!bucket_meta.valid) { + cuda::std::uint32_t bucket_key; + if (not load(bitmap, byte_offset, bucket_key)) { return; } + byte_offset += sizeof(cuda::std::uint32_t); + + roaring_bitmap_metadata bucket_meta{bitmap.subview(byte_offset)}; + if (not bucket_meta.valid) { + valid = false; + return; + } + if (bucket_meta.num_keys > cuda::std::numeric_limits::max() - num_keys or + bucket_meta.size_bytes > + cuda::std::numeric_limits::max() - byte_offset) { valid = false; + NV_IF_TARGET( + NV_IS_HOST, + CUCO_FAIL("Invalid bitmap format: bitmap size overflow");) // TODO device error handling return; } num_keys += bucket_meta.num_keys; - byte_offset += bucket_meta.size_bytes; // skip bucket + byte_offset += bucket_meta.size_bytes; } size_bytes = byte_offset; valid = true; diff --git a/include/cuco/roaring_bitmap.cuh b/include/cuco/roaring_bitmap.cuh index d01b990f4..282ac899c 100644 --- a/include/cuco/roaring_bitmap.cuh +++ b/include/cuco/roaring_bitmap.cuh @@ -10,6 +10,7 @@ #include #include +#include #include namespace cuco::experimental { @@ -41,6 +42,23 @@ class roaring_bitmap { * @brief Constructs a `roaring_bitmap` by copying the serialized bytes to device-accessible * storage. * + * @param bitmap Serialized bitmap bytes in host memory + * @param alloc Allocator used to allocate device-accessible storage + * @param stream CUDA stream used for device memory operations during construction + */ + roaring_bitmap(cuda::std::span bitmap, + Allocator const& alloc = {}, + cuda::stream_ref stream = cuda::stream_ref{cudaStream_t{nullptr}}); + + /** + * @brief Constructs a `roaring_bitmap` by copying the serialized bytes to device-accessible + * storage. + * + * @deprecated Use the overload accepting `cuda::std::span` to enable bounds validation. + * + * @note The caller must ensure `bitmap` points to a complete, valid serialized bitmap. This + * overload cannot validate the bounds of the serialized data. + * * @param bitmap Pointer to the beginning of the serialized bitmap in host memory * @param alloc Allocator used to allocate device-accessible storage * @param stream CUDA stream used for device memory operations during construction diff --git a/tests/roaring_bitmap/contains_test.cu b/tests/roaring_bitmap/contains_test.cu index 642a1da2c..41313167c 100644 --- a/tests/roaring_bitmap/contains_test.cu +++ b/tests/roaring_bitmap/contains_test.cu @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -70,7 +71,7 @@ bool check(std::string const& bitmap_file_path) file.close(); cuco::experimental::roaring_bitmap roaring_bitmap( - thrust::raw_pointer_cast(buffer.data())); + cuda::std::span{thrust::raw_pointer_cast(buffer.data()), buffer.size()}); auto keys = generate_keys(); thrust::device_vector contained(keys.size(), false); @@ -115,6 +116,51 @@ std::vector make_run_container_no_offsets_bitmap() } } // namespace +TEST_CASE("roaring_bitmap rejects truncated serialized data", "[roaring_bitmap]") +{ + auto const bytes = std::vector{ + cuda::std::byte{0x3B}, + cuda::std::byte{0x30}, + cuda::std::byte{0xFF}, + cuda::std::byte{0x00}}; // Run-container cookie declaring 256 containers + auto const bitmap = cuda::std::span{bytes.data(), bytes.size()}; + + REQUIRE_THROWS(cuco::experimental::roaring_bitmap{bitmap}); +} + +TEST_CASE("roaring_bitmap rejects out-of-bounds container offsets", "[roaring_bitmap]") +{ + auto const bytes = + std::vector{cuda::std::byte{0x3A}, + cuda::std::byte{0x30}, + cuda::std::byte{0x00}, + cuda::std::byte{0x00}, // No-run-container cookie + cuda::std::byte{0x01}, + cuda::std::byte{0x00}, + cuda::std::byte{0x00}, + cuda::std::byte{0x00}, // One container + cuda::std::byte{0x00}, + cuda::std::byte{0x00}, // Key + cuda::std::byte{0x00}, + cuda::std::byte{0x00}, // Cardinality minus one + cuda::std::byte{0xFF}, + cuda::std::byte{0xFF}, + cuda::std::byte{0xFF}, + cuda::std::byte{0xFF}}; // Container offset outside the buffer + auto const bitmap = cuda::std::span{bytes.data(), bytes.size()}; + + REQUIRE_THROWS(cuco::experimental::roaring_bitmap{bitmap}); +} + +TEST_CASE("64-bit roaring_bitmap rejects unbounded bucket counts", "[roaring_bitmap]") +{ + auto const bytes = + std::vector(sizeof(cuda::std::uint64_t), cuda::std::byte{0xFF}); + auto const bitmap = cuda::std::span{bytes.data(), bytes.size()}; + + REQUIRE_THROWS(cuco::experimental::roaring_bitmap{bitmap}); +} + TEST_CASE("roaring_bitmap run container without offsets", "[roaring_bitmap]") { // When run containers are present and the bitmap has fewer than 4 containers, the @@ -125,7 +171,7 @@ TEST_CASE("roaring_bitmap run container without offsets", "[roaring_bitmap]") std::memcpy(thrust::raw_pointer_cast(buffer.data()), bytes.data(), bytes.size()); cuco::experimental::roaring_bitmap roaring_bitmap( - thrust::raw_pointer_cast(buffer.data())); + cuda::std::span{thrust::raw_pointer_cast(buffer.data()), buffer.size()}); thrust::device_vector keys{1, 2, 3, 4}; thrust::device_vector contained(keys.size(), false);