diff --git a/cpp/include/cugraph/dynamic/memory_manager/bit_tree.hpp b/cpp/include/cugraph/dynamic/memory_manager/bit_tree.hpp new file mode 100644 index 0000000000..c710db457c --- /dev/null +++ b/cpp/include/cugraph/dynamic/memory_manager/bit_tree.hpp @@ -0,0 +1,161 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include + +#include +#include +#include +#include + +namespace cugraph { +namespace dynamic { + +/** + * @brief Tracks which fixed-size blocks are free in a block array. + * + * A block array is a large fixed-capacity buffer split into equally-sized blocks. This host-side + * bitmap records which blocks are available so insert() can hand out a free block quickly and + * remove() can return one. + * + * Based on Hornet's Vec-Tree bitmap. + */ +class bit_tree_t { + public: + bit_tree_t(size_t elements_per_block, size_t num_blocks) + { + CUGRAPH_EXPECTS(std::has_single_bit(elements_per_block), + "Invalid input argument: elements_per_block must be a power of 2."); + CUGRAPH_EXPECTS(std::has_single_bit(num_blocks), + "Invalid input argument: num_blocks must be a power of 2."); + + elements_per_block_ = elements_per_block; + num_blocks_ = num_blocks; + + log_elements_per_block_ = std::bit_width(static_cast(elements_per_block)) - size_t{1}; + auto log_num_blocks = std::bit_width(num_blocks) - size_t{1}; + auto log_fanout = std::bit_width(packed_bools_per_word()) - size_t{1}; + + auto ceil_log_fanout = (log_num_blocks + log_fanout - size_t{1}) / log_fanout; + num_internal_levels_ = + ((ceil_log_fanout) >= size_t{1}) ? (ceil_log_fanout - size_t{1}) : size_t{0}; + + // Internal summary bits cover the fanout tree above the leaf level. For L internal levels + // and fanout F, the full tree has (F^(L+1) - 1) / (F - 1) nodes; drop the root summary bit. + // e.g. 0 internal levels: 0, 1 internal levels: 32, 2 internal levels: 32*32 + 32, 3 internal + // levels: 32*32*32 + 32*32 + 32 + auto constexpr fanout = packed_bools_per_word(); + auto fanout_power = size_t{1}; + for (size_t i = 0; i <= num_internal_levels_; ++i) { + fanout_power *= fanout; + } + internal_summary_bits_ = ((fanout_power - size_t{1}) / (fanout - size_t{1})) - size_t{1}; + internal_words_ = packed_bool_size(internal_summary_bits_); + leaf_words_ = packed_bool_size(num_blocks_); + array_.assign(internal_words_ + leaf_words_, packed_bool_full_mask()); + } + + bit_tree_t() = delete; + + size_t insert() + { + CUGRAPH_EXPECTS(num_occupied_blocks_ < num_blocks_, "Cannot insert into a full bit tree."); + ++num_occupied_blocks_; + + size_t index = 0; + for (size_t i = 0; i < num_internal_levels_; ++i) { + auto pos = static_cast(std::countr_zero(array_[index / packed_bools_per_word()])); + index = (index + pos + size_t{1}) * packed_bools_per_word(); + } + index += static_cast(std::countr_zero(array_[index / packed_bools_per_word()])); + + clear_available_bit(index); + if (array_[index / packed_bools_per_word()] == packed_bool_empty_mask()) { + // Walk up internal levels, clearing parent bits when a whole child group is full. + for (size_t parent_index = index / packed_bools_per_word(); parent_index != size_t{0}; + parent_index /= packed_bools_per_word()) { + --parent_index; + clear_available_bit(parent_index); + if (array_[parent_index / packed_bools_per_word()] != packed_bool_empty_mask()) { break; } + } + } + + auto block_index = index - internal_summary_bits_; + return block_index; + } + + void remove(size_t block_index) + { + CUGRAPH_EXPECTS(num_occupied_blocks_ != 0, "Cannot remove from an empty bit tree."); + --num_occupied_blocks_; + + CUGRAPH_EXPECTS(block_index < num_blocks_, + "Invalid input argument: block_index is out of range."); + CUGRAPH_EXPECTS(is_available_bit_clear(last_level_ptr(), block_index), + "Invalid input argument: block is not allocated."); + + set_available_bit(last_level_ptr(), block_index); + block_index += internal_summary_bits_; + + // Walk up internal levels, setting parent bits until a parent already marks availability. + for (size_t parent_index = block_index / packed_bools_per_word(); parent_index != size_t{0}; + parent_index /= packed_bools_per_word()) { + --parent_index; + bool parent_already_available = + (array_[parent_index / packed_bools_per_word()] != packed_bool_empty_mask()); + set_available_bit(parent_index); + if (parent_already_available) { break; } + } + } + + size_t num_occupied_blocks() const { return num_occupied_blocks_; } + + bool full() const { return num_occupied_blocks_ == num_blocks_; } + + size_t num_elements_per_block() const { return elements_per_block_; } + + size_t num_blocks() const { return num_blocks_; } + + size_t log_elements_per_block() const { return log_elements_per_block_; } + + private: + uint32_t* last_level_ptr() { return array_.data() + internal_words_; } + + uint32_t const* last_level_ptr() const { return array_.data() + internal_words_; } + + void clear_available_bit(size_t bit_index) + { + array_[packed_bool_offset(bit_index)] &= ~packed_bool_mask(bit_index); + } + + static void set_available_bit(uint32_t* array, size_t bit_index) + { + array[packed_bool_offset(bit_index)] |= packed_bool_mask(bit_index); + } + + void set_available_bit(size_t bit_index) { set_available_bit(array_.data(), bit_index); } + + static bool is_available_bit_clear(uint32_t const* array, size_t bit_index) + { + return (array[packed_bool_offset(bit_index)] & packed_bool_mask(bit_index)) == + packed_bool_empty_mask(); + } + + size_t elements_per_block_{0}; + size_t log_elements_per_block_{0}; + size_t num_blocks_{0}; + size_t num_internal_levels_{0}; + size_t internal_summary_bits_{0}; + size_t internal_words_{0}; + size_t leaf_words_{0}; + + std::vector array_{}; + size_t num_occupied_blocks_{0}; +}; + +} // namespace dynamic +} // namespace cugraph diff --git a/cpp/include/cugraph/dynamic/memory_manager/block_array.hpp b/cpp/include/cugraph/dynamic/memory_manager/block_array.hpp new file mode 100644 index 0000000000..8e5fbf1345 --- /dev/null +++ b/cpp/include/cugraph/dynamic/memory_manager/block_array.hpp @@ -0,0 +1,148 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include + +#include + +#include + +#include +#include + +namespace cugraph { +namespace dynamic { + +/** + * @brief Block array of fixed-size blocks with a host-side bit tree to track available blocks. + * + * @tparam T Arithmetic type or cuda::std::tuple of arithmetic types. + */ +template +class block_array_t { + public: + static_assert(is_arithmetic_or_thrust_tuple_of_arithmetic_v, + "T must be an arithmetic type or cuda::std::tuple of arithmetic types."); + + using buffer_type = dataframe_buffer_type_t; + + block_array_t(size_t elements_per_block, size_t num_blocks, rmm::cuda_stream_view stream) + : bit_tree_(elements_per_block, num_blocks), + block_storage_(allocate_dataframe_buffer(num_blocks * elements_per_block, stream)) + { + } + + block_array_t() = delete; + block_array_t(block_array_t const&) = delete; + block_array_t& operator=(block_array_t const&) = delete; + + block_array_t(block_array_t&&) = default; + block_array_t& operator=(block_array_t&&) = default; + + template , int> = 0> + auto column_data() + { + return block_storage_.data(); + } + + template , int> = 0> + auto column_data() const + { + return block_storage_.data(); + } + + template , int> = 0> + auto column_data() + { + static_assert(I < thrust_tuple_size_or_one::value); + return std::get(block_storage_).data(); + } + + template , int> = 0> + auto column_data() const + { + static_assert(I < thrust_tuple_size_or_one::value); + return std::get(block_storage_).data(); + } + + template , int> = 0> + auto block_data(size_t block_index) + { + return column_data() + block_index * num_elements_per_block(); + } + + template , int> = 0> + auto block_data(size_t block_index) const + { + return column_data() + block_index * num_elements_per_block(); + } + + template , int> = 0> + auto block_data(size_t block_index) + { + return column_data() + block_index * num_elements_per_block(); + } + + template , int> = 0> + auto block_data(size_t block_index) const + { + return column_data() + block_index * num_elements_per_block(); + } + + template , int> = 0> + auto element_data(size_t element_index) + { + return column_data() + element_index; + } + + template , int> = 0> + auto element_data(size_t element_index) const + { + return column_data() + element_index; + } + + template , int> = 0> + auto element_data(size_t element_index) + { + return column_data() + element_index; + } + + template , int> = 0> + auto element_data(size_t element_index) const + { + return column_data() + element_index; + } + + size_t insert() { return bit_tree_.insert(); } + + void remove(size_t block_index) { bit_tree_.remove(block_index); } + + size_t num_elements_per_block() const { return bit_tree_.num_elements_per_block(); } + + size_t num_blocks() const { return bit_tree_.num_blocks(); } + + size_t num_elements() const { return size_dataframe_buffer(block_storage_); } + + bool full() const { return bit_tree_.full(); } + + std::byte const* block_array_key() const + { + if constexpr (is_thrust_tuple_v) { + return reinterpret_cast(column_data<0>()); + } else { + return reinterpret_cast(column_data()); + } + } + + private: + bit_tree_t bit_tree_; + buffer_type block_storage_; +}; + +} // namespace dynamic +} // namespace cugraph diff --git a/cpp/include/cugraph/dynamic/memory_manager/block_array_manager.hpp b/cpp/include/cugraph/dynamic/memory_manager/block_array_manager.hpp new file mode 100644 index 0000000000..a3c7ce2ba6 --- /dev/null +++ b/cpp/include/cugraph/dynamic/memory_manager/block_array_manager.hpp @@ -0,0 +1,109 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace cugraph { +namespace dynamic { + +/** + * @brief Handle returned by block_array_manager_t::insert() to locate an allocated block. + */ +struct block_access_data_t { + std::byte const* block_array_key{nullptr}; + size_t block_index{0}; + size_t num_elements_per_block{0}; +}; + +/** + * @brief binning memory manager of block arrays (with # elements per block: 2^0, 2^1, ..., 2^n) + * + * @tparam T Arithmetic type or cuda::std::tuple of arithmetic types. + * @param max_elements_per_block_array Maximum number of elements per block array. The maximum block + * size can't exceed this value. + */ +template +class block_array_manager_t { + public: + static_assert(is_arithmetic_or_thrust_tuple_of_arithmetic_v, + "T must be an arithmetic type or cuda::std::tuple of arithmetic types."); + + static constexpr size_t num_bins_v = sizeof(size_t) * 8; + + explicit block_array_manager_t(size_t max_elements_per_block_array = (size_t{1} << 23)) + : max_elements_per_block_array_(std::bit_ceil(max_elements_per_block_array)) + { + CUGRAPH_EXPECTS(std::has_single_bit(max_elements_per_block_array), + "Invalid input argument: max_elements_per_block_array must be a power of 2."); + } + + block_array_manager_t(block_array_manager_t const&) = delete; + block_array_manager_t& operator=(block_array_manager_t const&) = delete; + + block_array_manager_t(block_array_manager_t&&) = default; + block_array_manager_t& operator=(block_array_manager_t&&) = default; + + block_access_data_t insert(size_t num_elements_per_block, rmm::cuda_stream_view stream) + { + CUGRAPH_EXPECTS( + num_elements_per_block <= max_elements_per_block_array_, + "num_elements_per_block must be less than or equal to max_elements_per_block_array_."); + if (num_elements_per_block == size_t{0}) { return block_access_data_t{}; } + + auto const bin_index = std::bit_width(num_elements_per_block - size_t{1}); + for (auto& entry : block_arrays_[bin_index]) { + if (!entry.second.full()) { + auto block_index = entry.second.insert(); + return {entry.second.block_array_key(), block_index, entry.second.num_elements_per_block()}; + } + } + + size_t const elements_per_block = std::bit_ceil(num_elements_per_block); + + assert((max_elements_per_block_array_ % elements_per_block) == size_t{0}); + size_t const num_blocks = max_elements_per_block_array_ / elements_per_block; + + block_array_t new_block_array(elements_per_block, num_blocks, stream); + size_t block_index = new_block_array.insert(); + + auto block_array_key = new_block_array.block_array_key(); + block_arrays_[bin_index].emplace(block_array_key, std::move(new_block_array)); + return {block_array_key, block_index, elements_per_block}; + } + + void remove(block_access_data_t const& access_data) + { + CUGRAPH_EXPECTS(access_data.num_elements_per_block > size_t{0}, + "num_elements_per_block must be greater than 0."); + auto const bin_index = std::bit_width(access_data.num_elements_per_block - size_t{1}); + block_arrays_[bin_index].at(access_data.block_array_key).remove(access_data.block_index); + } + + void remove_all() noexcept + { + for (auto& bin : block_arrays_) { + bin.clear(); + } + } + + private: + std::array>, + static_cast(sizeof(T) * 8)> + block_arrays_{}; + size_t max_elements_per_block_array_; +}; + +} // namespace dynamic +} // namespace cugraph diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 3a85492678..06283dade6 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -573,6 +573,15 @@ ConfigureTest(LOOKUP_SRC_DST_TEST lookup/lookup_src_dst_test.cpp) # - K-hop Neighbors tests ------------------------------------------------------------------------- ConfigureTest(K_HOP_NBRS_TEST traversal/k_hop_nbrs_test.cpp) +################################################################################################### +# - Dynamic graph memory manager tests ------------------------------------------------------------ +ConfigureTest(DYNAMIC_BIT_TREE_TEST dynamic/memory_manager/bit_tree_test.cpp) + +ConfigureTest(DYNAMIC_BLOCK_ARRAY_TEST dynamic/memory_manager/block_array_test.cpp) + +ConfigureTest(DYNAMIC_BLOCK_ARRAY_MANAGER_TEST + dynamic/memory_manager/block_array_manager_test.cpp) + ################################################################################################### # - install tests --------------------------------------------------------------------------------- diff --git a/cpp/tests/dynamic/memory_manager/bit_tree_test.cpp b/cpp/tests/dynamic/memory_manager/bit_tree_test.cpp new file mode 100644 index 0000000000..a7cc8fbfb9 --- /dev/null +++ b/cpp/tests/dynamic/memory_manager/bit_tree_test.cpp @@ -0,0 +1,71 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "utilities/base_fixture.hpp" + +#include + +#include + +#include +#include +#include + +class Tests_BitTree : public ::testing::Test { + public: + Tests_BitTree() {} + + static void SetUpTestCase() {} + static void TearDownTestCase() {} + + virtual void SetUp() {} + virtual void TearDown() {} + + void run_insert_until_full_test() + { + cugraph::dynamic::bit_tree_t tree(size_t{4}, size_t{4}); + + EXPECT_EQ(tree.log_elements_per_block(), size_t{2}); + EXPECT_EQ(tree.num_occupied_blocks(), size_t{0}); + EXPECT_FALSE(tree.full()); + + std::vector slots; + slots.reserve(4); + for (size_t i = 0; i < size_t{4}; ++i) { + auto const slot = tree.insert(); + EXPECT_EQ(std::find(slots.begin(), slots.end(), slot), slots.end()); + slots.push_back(slot); + EXPECT_EQ(tree.num_occupied_blocks(), i + 1); + } + + EXPECT_TRUE(tree.full()); + EXPECT_EQ(tree.num_occupied_blocks(), size_t{4}); + } + + void run_remove_and_reuse_test() + { + cugraph::dynamic::bit_tree_t tree(size_t{4}, size_t{4}); + + std::vector slots; + slots.reserve(4); + for (size_t i = 0; i < size_t{4}; ++i) { + slots.push_back(tree.insert()); + } + + tree.remove(slots[1]); + EXPECT_FALSE(tree.full()); + EXPECT_EQ(tree.num_occupied_blocks(), size_t{3}); + + auto const reused = tree.insert(); + EXPECT_EQ(reused, slots[1]); + EXPECT_TRUE(tree.full()); + } +}; + +TEST_F(Tests_BitTree, InsertUntilFull) { run_insert_until_full_test(); } + +TEST_F(Tests_BitTree, RemoveAndReuse) { run_remove_and_reuse_test(); } + +CUGRAPH_TEST_PROGRAM_MAIN() diff --git a/cpp/tests/dynamic/memory_manager/block_array_manager_test.cpp b/cpp/tests/dynamic/memory_manager/block_array_manager_test.cpp new file mode 100644 index 0000000000..015b11b89b --- /dev/null +++ b/cpp/tests/dynamic/memory_manager/block_array_manager_test.cpp @@ -0,0 +1,170 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "utilities/base_fixture.hpp" + +#include +#include + +#include + +#include + +#include + +#include +#include + +class Tests_BlockArrayManager : public ::testing::Test { + public: + static constexpr size_t k_max_elements_per_block_array = size_t{16}; + + Tests_BlockArrayManager() {} + + static void SetUpTestCase() {} + static void TearDownTestCase() {} + + virtual void SetUp() {} + virtual void TearDown() {} +}; + +TEST_F(Tests_BlockArrayManager, InsertDegreeZero) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_manager_t manager( + Tests_BlockArrayManager::k_max_elements_per_block_array); + auto const access = manager.insert(size_t{0}, handle.get_stream()); + + EXPECT_EQ(access.block_array_key, nullptr); + EXPECT_EQ(access.block_index, size_t{0}); + EXPECT_EQ(access.num_elements_per_block, size_t{0}); +} + +TEST_F(Tests_BlockArrayManager, InsertReusesBlockArray) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_manager_t manager( + Tests_BlockArrayManager::k_max_elements_per_block_array); + + auto const first = manager.insert(size_t{4}, handle.get_stream()); + auto const second = manager.insert(size_t{4}, handle.get_stream()); + + EXPECT_NE(first.block_array_key, nullptr); + EXPECT_EQ(first.block_array_key, second.block_array_key); + EXPECT_NE(first.block_index, second.block_index); + EXPECT_EQ(first.num_elements_per_block, size_t{4}); + EXPECT_EQ(second.num_elements_per_block, size_t{4}); +} + +TEST_F(Tests_BlockArrayManager, InsertRoundsUpElementsPerBlock) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_manager_t manager( + Tests_BlockArrayManager::k_max_elements_per_block_array); + + auto const access = manager.insert(size_t{5}, handle.get_stream()); + + EXPECT_NE(access.block_array_key, nullptr); + EXPECT_EQ(access.block_index, size_t{0}); + EXPECT_EQ(access.num_elements_per_block, size_t{8}); +} + +TEST_F(Tests_BlockArrayManager, InsertDegreeOneUsesBinZero) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_manager_t manager( + Tests_BlockArrayManager::k_max_elements_per_block_array); + + auto const access = manager.insert(size_t{1}, handle.get_stream()); + + EXPECT_NE(access.block_array_key, nullptr); + EXPECT_EQ(access.block_index, size_t{0}); + EXPECT_EQ(access.num_elements_per_block, size_t{1}); +} + +TEST_F(Tests_BlockArrayManager, InsertExceedsMaxElementsPerBlockArray) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_manager_t manager( + Tests_BlockArrayManager::k_max_elements_per_block_array); + + EXPECT_THROW(manager.insert(size_t{32}, handle.get_stream()), cugraph::logic_error); +} + +TEST_F(Tests_BlockArrayManager, RemoveAndInsertOnSameBlockArray) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_manager_t manager( + Tests_BlockArrayManager::k_max_elements_per_block_array); + + auto const first = manager.insert(size_t{4}, handle.get_stream()); + auto const second = manager.insert(size_t{4}, handle.get_stream()); + EXPECT_NE(first.block_index, second.block_index); + + manager.remove(first); + + auto const third = manager.insert(size_t{4}, handle.get_stream()); + EXPECT_EQ(third.block_array_key, first.block_array_key); + EXPECT_EQ(third.num_elements_per_block, first.num_elements_per_block); + EXPECT_NE(third.block_index, second.block_index); +} + +TEST_F(Tests_BlockArrayManager, DifferentBinsUseDifferentBlockArrays) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_manager_t manager( + Tests_BlockArrayManager::k_max_elements_per_block_array); + + auto const small = manager.insert(size_t{4}, handle.get_stream()); + auto const large = manager.insert(size_t{8}, handle.get_stream()); + + EXPECT_NE(small.block_array_key, large.block_array_key); + EXPECT_EQ(small.num_elements_per_block, size_t{4}); + EXPECT_EQ(large.num_elements_per_block, size_t{8}); +} + +TEST_F(Tests_BlockArrayManager, RemoveAllClearsAllocations) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_manager_t manager( + Tests_BlockArrayManager::k_max_elements_per_block_array); + manager.insert(size_t{4}, handle.get_stream()); + manager.remove_all(); + + auto const access = manager.insert(size_t{4}, handle.get_stream()); + EXPECT_NE(access.block_array_key, nullptr); + EXPECT_EQ(access.num_elements_per_block, size_t{4}); +} + +TEST_F(Tests_BlockArrayManager, MultiColumnInsertAndRemove) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_manager_t> manager( + Tests_BlockArrayManager::k_max_elements_per_block_array); + + auto const first = manager.insert(size_t{4}, handle.get_stream()); + auto const second = manager.insert(size_t{4}, handle.get_stream()); + + EXPECT_NE(first.block_array_key, nullptr); + EXPECT_EQ(first.block_array_key, second.block_array_key); + EXPECT_NE(first.block_index, second.block_index); + + manager.remove(first); + + auto const third = manager.insert(size_t{4}, handle.get_stream()); + EXPECT_EQ(third.block_array_key, first.block_array_key); + EXPECT_NE(third.block_index, second.block_index); +} + +CUGRAPH_TEST_PROGRAM_MAIN() diff --git a/cpp/tests/dynamic/memory_manager/block_array_test.cpp b/cpp/tests/dynamic/memory_manager/block_array_test.cpp new file mode 100644 index 0000000000..3aacea25fc --- /dev/null +++ b/cpp/tests/dynamic/memory_manager/block_array_test.cpp @@ -0,0 +1,204 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "utilities/base_fixture.hpp" +#include "utilities/conversion_utilities.hpp" + +#include + +#include + +#include + +#include + +#include +#include +#include +#include +#include + +class Tests_BlockArray : public ::testing::Test { + public: + static constexpr size_t k_elements_per_block = 4; + static constexpr size_t k_num_blocks = 4; + static constexpr size_t k_num_elements = k_elements_per_block * k_num_blocks; + + Tests_BlockArray() {} + + static void SetUpTestCase() {} + static void TearDownTestCase() {} + + virtual void SetUp() {} + virtual void TearDown() {} + + void run_insert_until_full_test(raft::handle_t const& handle) + { + cugraph::dynamic::block_array_t block_array( + k_elements_per_block, k_num_blocks, handle.get_stream()); + + EXPECT_NE(block_array.block_array_key(), nullptr); + EXPECT_EQ(block_array.num_elements_per_block(), k_elements_per_block); + EXPECT_EQ(block_array.num_blocks(), k_num_blocks); + EXPECT_EQ(block_array.num_elements(), k_num_elements); + EXPECT_FALSE(block_array.full()); + + std::vector block_indices; + block_indices.reserve(k_num_blocks); + for (size_t i = 0; i < k_num_blocks; ++i) { + auto const block_index = block_array.insert(); + EXPECT_EQ(std::find(block_indices.begin(), block_indices.end(), block_index), + block_indices.end()); + block_indices.push_back(block_index); + EXPECT_EQ(block_index, i); + } + + EXPECT_TRUE(block_array.full()); + } + + void run_remove_and_reuse_test(raft::handle_t const& handle) + { + cugraph::dynamic::block_array_t block_array( + k_elements_per_block, k_num_blocks, handle.get_stream()); + + std::vector block_indices; + block_indices.reserve(k_num_blocks); + for (size_t i = 0; i < k_num_blocks; ++i) { + block_indices.push_back(block_array.insert()); + } + + block_array.remove(block_indices[1]); + EXPECT_FALSE(block_array.full()); + + auto const reused = block_array.insert(); + EXPECT_EQ(reused, block_indices[1]); + EXPECT_TRUE(block_array.full()); + } +}; + +TEST_F(Tests_BlockArray, InsertUntilFull) +{ + raft::handle_t handle; + run_insert_until_full_test(handle); +} + +TEST_F(Tests_BlockArray, RemoveAndReuse) +{ + raft::handle_t handle; + run_remove_and_reuse_test(handle); +} + +TEST_F(Tests_BlockArray, BlockAndElementDataPointers) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_t block_array( + Tests_BlockArray::k_elements_per_block, Tests_BlockArray::k_num_blocks, handle.get_stream()); + + auto const block_index = size_t{2}; + auto const in_block_offset = size_t{3}; + auto const element_index = block_index * Tests_BlockArray::k_elements_per_block + in_block_offset; + + EXPECT_EQ(block_array.block_data(block_index), + block_array.column_data() + block_index * Tests_BlockArray::k_elements_per_block); + EXPECT_EQ(block_array.element_data(element_index), block_array.column_data() + element_index); + EXPECT_EQ(block_array.block_data(block_index) + in_block_offset, + block_array.element_data(element_index)); +} + +TEST_F(Tests_BlockArray, SingleColumnReadWrite) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_t block_array( + Tests_BlockArray::k_elements_per_block, Tests_BlockArray::k_num_blocks, handle.get_stream()); + + std::vector expected(Tests_BlockArray::k_num_elements); + std::iota(expected.begin(), expected.end(), 0); + + raft::update_device( + block_array.column_data(), expected.data(), expected.size(), handle.get_stream()); + handle.sync_stream(); + + auto actual = cugraph::test::to_host( + handle, + raft::device_span(block_array.column_data(), Tests_BlockArray::k_num_elements)); + + EXPECT_EQ(expected, actual); +} + +TEST_F(Tests_BlockArray, BlockColumnReadWrite) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_t block_array( + Tests_BlockArray::k_elements_per_block, Tests_BlockArray::k_num_blocks, handle.get_stream()); + + auto const block_index = block_array.insert(); + + std::vector expected(Tests_BlockArray::k_elements_per_block); + std::iota(expected.begin(), expected.end(), static_cast(block_index * 10)); + + raft::update_device( + block_array.block_data(block_index), expected.data(), expected.size(), handle.get_stream()); + handle.sync_stream(); + + auto actual = cugraph::test::to_host( + handle, + raft::device_span(block_array.block_data(block_index), + Tests_BlockArray::k_elements_per_block)); + + EXPECT_EQ(expected, actual); +} + +TEST_F(Tests_BlockArray, MultiColumnBuffer) +{ + raft::handle_t handle; + + constexpr size_t num_items = 4; + cugraph::dynamic::block_array_t> block_array( + Tests_BlockArray::k_elements_per_block, size_t{1}, handle.get_stream()); + + std::vector int_values(num_items); + std::iota(int_values.begin(), int_values.end(), 10); + + std::vector float_values(num_items); + for (size_t i = 0; i < num_items; ++i) { + float_values[i] = static_cast(i) * 0.5f; + } + + raft::update_device( + block_array.column_data<0>(), int_values.data(), int_values.size(), handle.get_stream()); + raft::update_device( + block_array.column_data<1>(), float_values.data(), float_values.size(), handle.get_stream()); + handle.sync_stream(); + + auto actual_int = cugraph::test::to_host( + handle, raft::device_span(block_array.column_data<0>(), num_items)); + auto actual_float = cugraph::test::to_host( + handle, raft::device_span(block_array.column_data<1>(), num_items)); + + EXPECT_EQ(int_values, actual_int); + EXPECT_EQ(float_values, actual_float); +} + +TEST_F(Tests_BlockArray, MultiColumnBlockDataPointers) +{ + raft::handle_t handle; + + cugraph::dynamic::block_array_t> block_array( + Tests_BlockArray::k_elements_per_block, size_t{1}, handle.get_stream()); + + EXPECT_EQ(block_array.block_data<0>(0), block_array.column_data<0>()); + EXPECT_EQ(block_array.block_data<1>(0), block_array.column_data<1>()); + EXPECT_EQ(block_array.element_data<0>(2), block_array.column_data<0>() + 2); + EXPECT_EQ(block_array.element_data<1>(2), block_array.column_data<1>() + 2); + EXPECT_NE(block_array.block_array_key(), + reinterpret_cast(block_array.column_data<1>())); + EXPECT_EQ(block_array.block_array_key(), + reinterpret_cast(block_array.column_data<0>())); +} + +CUGRAPH_TEST_PROGRAM_MAIN()