kv_store.cuh

namespace cugraph

template<typename key_t, typename value_t, bool use_binary_search = true>
class kv_store_t

Public Types

using key_type = key_t

using value_type = value_t

Public Functions

inline kv_store_t(rmm::cuda_stream_view stream)

template<bool binary_search = use_binary_search>
inline kv_store_t(

size_t capacity,

key_t invalid_key,

value_t invalid_value,

rmm::cuda_stream_view stream,

std::enable_if_t<!binary_search, int32_t> = 0

)

template<typename KeyIterator, typename ValueIterator, bool binary_search = use_binary_search>
inline kv_store_t(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

value_t invalid_value,

bool key_sorted,

rmm::cuda_stream_view stream,

std::enable_if_t<binary_search, int32_t> = 0

)

template<typename KeyIterator, typename ValueIterator, bool binary_search = use_binary_search>
inline kv_store_t(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

key_t invalid_key,

value_t invalid_value,

rmm::cuda_stream_view stream,

std::enable_if_t<!binary_search, int32_t> = 0

)

template<bool binary_search = use_binary_search>
inline kv_store_t(

rmm::device_uvector<key_t> &&keys,

decltype(allocate_dataframe_buffer<value_t>(0, rmm::cuda_stream_view{})) &&values,

value_t invalid_value,

bool key_sorted,

rmm::cuda_stream_view stream,

std::enable_if_t<binary_search, int32_t> = 0

)

template<typename KeyIterator, typename ValueIterator, bool binary_search = use_binary_search>
inline std::enable_if_t<!binary_search, void> insert(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

rmm::cuda_stream_view stream

)

template<typename KeyIterator, typename ValueIterator, typename StencilIterator, typename PredOp, bool binary_search = use_binary_search>
inline std::enable_if_t<!binary_search, void> insert_if(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

StencilIterator stencil_first,

PredOp pred_op,

rmm::cuda_stream_view stream

)

template<typename KeyIterator, typename ValueIterator, bool binary_search = use_binary_search>
inline std::enable_if_t<!binary_search, void> insert_and_assign(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

rmm::cuda_stream_view stream

)

template<typename KeyIterator, typename ValueIterator, typename StencilIterator, typename PredOp, bool binary_search = use_binary_search>
inline std::enable_if_t<!binary_search, void> insert_and_assign_if(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

StencilIterator stencil_first,

PredOp pred_op,

rmm::cuda_stream_view stream

)

inline auto retrieve_all(rmm::cuda_stream_view stream) const

inline auto release(rmm::cuda_stream_view stream)

inline auto view() const

template<bool binary_search = use_binary_search>
inline std::enable_if_t<!binary_search, key_t> invalid_key(

) const

inline value_t invalid_value() const

inline size_t size() const

inline size_t capacity() const

Private Members

std::conditional_t<use_binary_search, detail::kv_binary_search_store_t<key_t, value_t>, detail::kv_cuco_store_t<key_t, value_t>> store_

namespace detail

template<typename KeyIterator>
struct kv_binary_search_contains_op_t

Public Types

using key_type = typename thrust::iterator_traits<KeyIterator>::value_type

Public Functions

inline __device__ bool operator() (key_type key) const

Public Members

KeyIterator store_key_first = {}

KeyIterator store_key_last = {}

template<typename KeyIterator, typename ValueIterator>
struct kv_binary_search_find_op_t

Public Types

using key_type = typename thrust::iterator_traits<KeyIterator>::value_type

using value_type = typename thrust::iterator_traits<ValueIterator>::value_type

Public Functions

inline __device__ value_type operator() (key_type key) const

Public Members

KeyIterator store_key_first = {}

KeyIterator store_key_last = {}

ValueIterator store_value_first = {}

value_type invalid_value = {}

template<typename ViewType>
struct kv_binary_search_store_device_view_t

Public Types

using key_type = typename ViewType::key_type

using value_type = typename ViewType::value_type

Public Functions

inline __host__ kv_binary_search_store_device_view_t(ViewType view)

inline __device__ value_type find (key_type key) const

Public Members

ViewType::key_iterator store_key_first = {}

ViewType::key_iterator store_key_last = {}

ViewType::value_iterator store_value_first = {}

value_type invalid_value = {}

template<typename key_t, typename value_t>
class kv_binary_search_store_t

Public Types

using key_type = key_t

using value_type = value_t

Public Functions

inline kv_binary_search_store_t(rmm::cuda_stream_view stream)

template<typename KeyIterator, typename ValueIterator>
inline kv_binary_search_store_t(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

value_t invalid_value,

bool key_sorted,

rmm::cuda_stream_view stream

)

inline kv_binary_search_store_t(

rmm::device_uvector<key_t> &&keys,

decltype(allocate_dataframe_buffer<value_t>(0, rmm::cuda_stream_view{})) &&values,

value_t invalid_value,

bool key_sorted,

rmm::cuda_stream_view stream

)

inline auto retrieve_all(rmm::cuda_stream_view stream)

inline auto release(rmm::cuda_stream_view stream)

inline key_t const *store_key_first() const

inline key_t const *store_key_last() const

inline auto store_value_first() const

inline value_t invalid_value() const

inline size_t size() const

inline size_t capacity() const

Private Members

rmm::device_uvector<key_t> store_keys_

decltype(allocate_dataframe_buffer<value_t>(0, rmm::cuda_stream_view{})) store_values_

value_t invalid_value_ = {}

template<typename KeyIterator, typename ValueIterator>
class kv_binary_search_store_view_t

Public Types

using key_type = std::remove_cv_t<typename thrust::iterator_traits<KeyIterator>::value_type>

using value_type = std::remove_cv_t<typename thrust::iterator_traits<ValueIterator>::value_type>

using key_iterator = KeyIterator

using value_iterator = ValueIterator

Public Functions

inline kv_binary_search_store_view_t(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

value_type invalid_value

)

template<typename QueryKeyIterator, typename ResultValueIterator>
inline void find(

QueryKeyIterator key_first,

QueryKeyIterator key_last,

ResultValueIterator value_first,

rmm::cuda_stream_view stream

) const

template<typename QueryKeyIterator, typename ResultValueIterator>
inline void contains(

QueryKeyIterator key_first,

QueryKeyIterator key_last,

ResultValueIterator value_first,

rmm::cuda_stream_view stream

) const

inline KeyIterator store_key_first() const

inline KeyIterator store_key_last() const

inline ValueIterator store_value_first() const

inline value_type invalid_value() const

Public Static Attributes

static bool binary_search = true

Private Members

KeyIterator store_key_first_ = {}

KeyIterator store_key_last_ = {}

ValueIterator store_value_first_ = {}

value_type invalid_value_ = {}

template<typename RefType, typename key_t, typename value_t>
struct kv_cuco_insert_and_assign_t

Public Functions

inline __device__ void operator() (cuda::std::tuple< key_t, value_t > pair)

Public Members

RefType device_ref = {}

template<typename RefType, typename KeyIterator>
struct kv_cuco_insert_and_increment_t

Public Functions

inline __device__ size_t operator() (size_t i)

Public Members

RefType device_ref = {}

KeyIterator key_first = {}

size_t *counter = {nullptr}

size_t invalid_idx = {}

template<typename RefType, typename KeyIterator, typename StencilIterator, typename PredOp>
struct kv_cuco_insert_if_and_increment_t

Public Functions

inline __device__ size_t operator() (size_t i)

Public Members

RefType device_ref = {}

KeyIterator key_first = {}

StencilIterator stencil_first = {}

PredOp pred_op = {}

size_t *counter = {nullptr}

size_t invalid_idx = {}

template<typename ViewType>
struct kv_cuco_store_find_device_view_t

Public Types

using key_type = typename ViewType::key_type

using value_type = typename ViewType::value_type

using cuco_store_device_ref_type = typename ViewType::cuco_map_type::ref_type<cuco::find_tag>

Public Functions

inline __host__ kv_cuco_store_find_device_view_t(ViewType view)

inline __device__ value_type find (key_type key) const

Public Members

cuco_store_device_ref_type cuco_store_device_ref = {}

std::conditional_t<!std::is_arithmetic_v<value_type>, typename ViewType::value_iterator, std::byte> store_value_first = {}

value_type invalid_value = {}

template<typename key_t, typename value_t>
class kv_cuco_store_t

Public Types

using key_type = key_t

using value_type = value_t

using value_buffer_type = decltype(allocate_dataframe_buffer<value_t>(0, rmm::cuda_stream_view{}))

using const_value_iterator = std::invoke_result_t<decltype(get_dataframe_buffer_cbegin<value_buffer_type>), value_buffer_type&>

using cuco_map_type = cuco::static_map<key_t, std::conditional_t<std::is_arithmetic_v<value_t>, value_t, size_t>, cuco::extent<std::size_t>, cuda::thread_scope_device, cuda::std::equal_to<key_t>, cuco::linear_probing<1, cuco::murmurhash3_32<key_t>>, rmm::mr::polymorphic_allocator<std::byte>, cuco_storage_type>

Public Functions

inline kv_cuco_store_t(rmm::cuda_stream_view stream)

inline kv_cuco_store_t(

size_t capacity,

key_t invalid_key,

value_t invalid_value,

rmm::cuda_stream_view stream

)

template<typename KeyIterator, typename ValueIterator>
inline kv_cuco_store_t(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

key_t invalid_key,

value_t invalid_value,

rmm::cuda_stream_view stream

)

template<typename KeyIterator, typename ValueIterator>
inline void insert(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

rmm::cuda_stream_view stream

)

template<typename KeyIterator, typename ValueIterator, typename StencilIterator, typename PredOp>
inline void insert_if(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

StencilIterator stencil_first,

PredOp pred_op,

rmm::cuda_stream_view stream

)

template<typename KeyIterator, typename ValueIterator>
inline void insert_and_assign(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

rmm::cuda_stream_view stream

)

template<typename KeyIterator, typename ValueIterator, typename StencilIterator, typename PredOp>
inline void insert_and_assign_if(

KeyIterator key_first,

KeyIterator key_last,

ValueIterator value_first,

StencilIterator stencil_first,

PredOp pred_op,

rmm::cuda_stream_view stream

)

inline auto retrieve_all(rmm::cuda_stream_view stream)

inline auto release(rmm::cuda_stream_view stream)

inline cuco_map_type const *cuco_store_ptr() const

template<typename type = value_t>
inline std::enable_if_t<!std::is_arithmetic_v<type>, const_value_iterator> store_value_first(

) const

inline key_t invalid_key() const

inline value_t invalid_value() const

inline size_t size() const

inline size_t capacity() const

Private Functions

inline void allocate(

size_t num_keys,

key_t invalid_key,

value_t invalid_value,

rmm::cuda_stream_view stream

)

Private Members

std::unique_ptr<cuco_map_type> cuco_store_ = {nullptr}

decltype(allocate_optional_dataframe_buffer<std::conditional_t<!std::is_arithmetic_v<value_t>, value_t, void>>(0, rmm::cuda_stream_view{})) store_values_

std::conditional_t<!std::is_arithmetic_v<value_t>, value_t, std::byte> invalid_value_ = {}

size_t capacity_ = {0}

size_t size_{0}

template<typename key_t, typename ValueIterator>
class kv_cuco_store_view_t

Public Types

using key_type = key_t

using value_type = std::remove_cv_t<typename thrust::iterator_traits<ValueIterator>::value_type>

using value_iterator = ValueIterator

using cuco_map_type = cuco::static_map<key_t, std::conditional_t<std::is_arithmetic_v<value_type>, value_type, size_t>, cuco::extent<std::size_t>, cuda::thread_scope_device, cuda::std::equal_to<key_t>, cuco::linear_probing<1, cuco::murmurhash3_32<key_t>>, rmm::mr::polymorphic_allocator<std::byte>, cuco_storage_type>

Public Functions

template<typename type = value_type>
inline kv_cuco_store_view_t(

cuco_map_type const *store,

std::enable_if_t<std::is_arithmetic_v<type>, int32_t> = 0

)

template<typename type = value_type>
inline kv_cuco_store_view_t(

cuco_map_type const *store,

ValueIterator value_first,

type invalid_value,

std::enable_if_t<!std::is_arithmetic_v<type>, int32_t> = 0

)

template<typename QueryKeyIterator, typename ResultValueIterator>
inline void find(

QueryKeyIterator key_first,

QueryKeyIterator key_last,

ResultValueIterator value_first,

rmm::cuda_stream_view stream

) const

template<typename QueryKeyIterator, typename ResultValueIterator>
inline void contains(

QueryKeyIterator key_first,

QueryKeyIterator key_last,

ResultValueIterator value_first,

rmm::cuda_stream_view stream

) const

inline auto cuco_store_find_device_ref() const

template<typename type = value_type>
inline std::enable_if_t<!std::is_arithmetic_v<type>, ValueIterator> store_value_first(

) const

inline key_t invalid_key() const

inline value_type invalid_value() const

Public Static Attributes

static bool binary_search = false

Private Members

cuco_map_type const *cuco_store_ = {}

std::conditional_t<!std::is_arithmetic_v<value_type>, ValueIterator, std::byte> store_value_first_ = {}

std::conditional_t<!std::is_arithmetic_v<value_type>, value_type, std::byte> invalid_value_ = {}