diff --git a/CMakeLists.txt b/CMakeLists.txt index 22e8902b..c9f66145 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -108,6 +108,11 @@ target_compile_features(merlin_hashtable_test PUBLIC cxx_std_14) set_target_properties(merlin_hashtable_test PROPERTIES CUDA_ARCHITECTURES OFF) TARGET_LINK_LIBRARIES(merlin_hashtable_test gtest_main) +add_executable(merlin_hashtable_device_api_test tests/merlin_hashtable_device_api_test.cc.cu) +target_compile_features(merlin_hashtable_device_api_test PUBLIC cxx_std_14) +set_target_properties(merlin_hashtable_device_api_test PROPERTIES CUDA_ARCHITECTURES OFF) +TARGET_LINK_LIBRARIES(merlin_hashtable_device_api_test gtest_main) + add_executable(find_or_insert_test tests/find_or_insert_test.cc.cu) target_compile_features(find_or_insert_test PUBLIC cxx_std_14) set_target_properties(find_or_insert_test PROPERTIES CUDA_ARCHITECTURES OFF) diff --git a/include/BUILD b/include/BUILD index d4255a4f..7b7705b3 100644 --- a/include/BUILD +++ b/include/BUILD @@ -17,7 +17,9 @@ cuda_cc_library( cuda_cc_library( name = "merlin_hashtable", hdrs = [ + "merlin_hashtable_base.hpp", "merlin_hashtable.cuh", + "merlin_hashtable_device.cuh", ], visibility = [ "//visibility:public", diff --git a/include/merlin/core_kernels/kernel_utils.cuh b/include/merlin/core_kernels/kernel_utils.cuh index 1b9b8e5c..8f377bee 100644 --- a/include/merlin/core_kernels/kernel_utils.cuh +++ b/include/merlin/core_kernels/kernel_utils.cuh @@ -642,6 +642,34 @@ __forceinline__ __device__ uint32_t get_start_position( return start_idx; } +// Read-only, no-tile-sync fast path lookup for static tables. +// Assumptions: +// - Table is immutable during lookup (static), keys are non-overlapping per +// query batch. +// - No atomic ops or tile-wide ballots; each thread scans positions with stride +// TILE_SIZE. +// - Whoever finds the key may directly write results; others can return early +// on EMPTY_KEY. +template +__device__ __forceinline__ int find_without_lock_readonly_no_sync( + Bucket* __restrict__ bucket, const K desired_key, + const uint32_t start_idx, const uint32_t bucket_max_size, + const uint32_t rank) { + const uint32_t mask = bucket_max_size - 1; + for (uint32_t tile_offset = 0; tile_offset < bucket_max_size; + tile_offset += TILE_SIZE) { + const uint32_t pos = (start_idx + rank + tile_offset) & mask; + const K current_key = *reinterpret_cast(bucket->keys(pos)); + if (current_key == desired_key) { + return static_cast(pos); + } + if (current_key == static_cast(EMPTY_KEY)) { + return -1; + } + } + return -1; +} + template __device__ __forceinline__ OccupyResult find_without_lock( cg::thread_block_tile g, Bucket* __restrict__ bucket, diff --git a/include/merlin_hashtable.cuh b/include/merlin_hashtable.cuh index e952ce30..fd0968c6 100644 --- a/include/merlin_hashtable.cuh +++ b/include/merlin_hashtable.cuh @@ -37,6 +37,7 @@ #include "merlin/multi_vector.hpp" #include "merlin/types.cuh" #include "merlin/utils.cuh" +#include "merlin_hashtable_base.hpp" namespace nv { namespace merlin { @@ -184,706 +185,6 @@ static constexpr auto& thrust_par = thrust::cuda::par_nosync; static constexpr auto& thrust_par = thrust::cuda::par; #endif -template -class HashTableBase { - public: - using size_type = size_t; - using key_type = K; - using value_type = V; - using score_type = S; - using allocator_type = BaseAllocator; - - public: - virtual ~HashTableBase() {} - - /** - * @brief Initialize a merlin::HashTable. - * - * @param options The configuration options. - */ - virtual void init(const HashTableOptions& options, - allocator_type* allocator = nullptr) = 0; - - /** - * @brief Insert new key-value-score tuples into the hash table. - * If the key already exists, the values and scores are assigned new values. - * - * If the target bucket is full, the keys with minimum score will be - * overwritten by new key unless the score of the new key is even less than - * minimum score of the target bucket. - * - * @param n Number of key-value-score tuples to insert or assign. - * @param keys The keys to insert on GPU-accessible memory with shape - * (n). - * @param values The values to insert on GPU-accessible memory with - * shape (n, DIM). - * @param scores The scores to insert on GPU-accessible memory with shape - * (n). - * @parblock - * The scores should be a `uint64_t` value for built-in strategies. For - * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. - * You can specify a value such as the timestamp of the key insertion or - * number of key occurrences to perform a customized eviction strategy. - * - * The @p scores should be `nullptr`, when the LRU eviction strategy is - * applied. - * @endparblock - * - * @param stream The CUDA stream that is used to execute the operation. - * @param unique_key If all keys in the same batch are unique. - * - * @param ignore_evict_strategy A boolean option indicating whether if - * the insert_or_assign ignores the evict strategy of table with current - * scores anyway. If true, it does not check whether the scores conforms to - * the evict strategy. If false, it requires the scores follow the evict - * strategy of table. - */ - virtual void insert_or_assign(const size_type n, - const key_type* keys, // (n) - const value_type* values, // (n, DIM) - const score_type* scores = nullptr, // (n) - cudaStream_t stream = 0, bool unique_key = true, - bool ignore_evict_strategy = false) = 0; - - /** - * @brief Insert new key-value-score tuples into the hash table. - * If the key already exists, the values and scores are assigned new values. - * - * If the target bucket is full, the keys with minimum score will be - * overwritten by new key unless the score of the new key is even less than - * minimum score of the target bucket. The overwritten key with minimum - * score will be evicted, with its values and score, to evicted_keys, - * evicted_values, evcted_scores seperately in compact format. - * - * @param n Number of key-value-score tuples to insert or assign. - * @param keys The keys to insert on GPU-accessible memory with shape - * (n). - * @param values The values to insert on GPU-accessible memory with - * shape (n, DIM). - * @param scores The scores to insert on GPU-accessible memory with shape - * (n). - * @param scores The scores to insert on GPU-accessible memory with shape - * (n). - * @params evicted_keys The output of keys replaced with minimum score. - * @params evicted_values The output of values replaced with minimum score on - * keys. - * @params evicted_scores The output of scores replaced with minimum score on - * keys. - * @parblock - * The scores should be a `uint64_t` value for built-in strategies. For - * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. - * You can specify a value such as the timestamp of the key insertion or - * number of key occurrences to perform a customized eviction strategy. - * - * The @p scores should be `nullptr`, when the LRU eviction strategy is - * applied. - * @endparblock - * - * @param d_evicted_counter The number of elements evicted on GPU-accessible - * memory. @notice The caller should guarantee it is set to `0` before - * calling. - * @param stream The CUDA stream that is used to execute the operation. - * @param unique_key If all keys in the same batch are unique. - * - * @param ignore_evict_strategy A boolean option indicating whether if - * the insert_or_assign ignores the evict strategy of table with current - * scores anyway. If true, it does not check whether the scores confroms to - * the evict strategy. If false, it requires the scores follow the evict - * strategy of table. - */ - virtual void insert_and_evict(const size_type n, - const key_type* keys, // (n) - const value_type* values, // (n, DIM) - const score_type* scores, // (n) - key_type* evicted_keys, // (n) - value_type* evicted_values, // (n, DIM) - score_type* evicted_scores, // (n) - size_type* d_evicted_counter, // (1) - cudaStream_t stream = 0, bool unique_key = true, - bool ignore_evict_strategy = false) = 0; - - /** - * @brief Insert new key-value-score tuples into the hash table. - * If the key already exists, the values and scores are assigned new values. - * - * If the target bucket is full, the keys with minimum score will be - * overwritten by new key unless the score of the new key is even less than - * minimum score of the target bucket. The overwritten key with minimum - * score will be evicted, with its values and score, to evicted_keys, - * evicted_values, evcted_scores seperately in compact format. - * - * @param n Number of key-value-score tuples to insert or assign. - * @param keys The keys to insert on GPU-accessible memory with shape - * (n). - * @param values The values to insert on GPU-accessible memory with - * shape (n, DIM). - * @param scores The scores to insert on GPU-accessible memory with shape - * (n). - * @param scores The scores to insert on GPU-accessible memory with shape - * (n). - * @params evicted_keys The output of keys replaced with minimum score. - * @params evicted_values The output of values replaced with minimum score on - * keys. - * @params evicted_scores The output of scores replaced with minimum score on - * keys. - * @parblock - * The scores should be a `uint64_t` value for built-in strategies. For - * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. - * You can specify a value such as the timestamp of the key insertion or - * number of key occurrences to perform a customized eviction strategy. - * - * The @p scores should be `nullptr`, when the LRU eviction strategy is - * applied. - * @endparblock - * - * @param stream The CUDA stream that is used to execute the operation. - * @param unique_key If all keys in the same batch are unique. - * - * @param ignore_evict_strategy A boolean option indicating whether if - * the insert_or_assign ignores the evict strategy of table with current - * scores anyway. If true, it does not check whether the scores confroms to - * the evict strategy. If false, it requires the scores follow the evict - * strategy of table. - * - * @return The number of elements evicted. - */ - virtual size_type insert_and_evict(const size_type n, - const key_type* keys, // (n) - const value_type* values, // (n, DIM) - const score_type* scores, // (n) - key_type* evicted_keys, // (n) - value_type* evicted_values, // (n, DIM) - score_type* evicted_scores, // (n) - cudaStream_t stream = 0, - bool unique_key = true, - bool ignore_evict_strategy = false) = 0; - - /** - * Searches for each key in @p keys in the hash table. - * If the key is found and the corresponding value in @p accum_or_assigns is - * `true`, the @p vectors_or_deltas is treated as a delta to the old - * value, and the delta is added to the old value of the key. - * - * If the key is not found and the corresponding value in @p accum_or_assigns - * is `false`, the @p vectors_or_deltas is treated as a new value and the - * key-value pair is updated in the table directly. - * - * @note When the key is found and the value of @p accum_or_assigns is - * `false`, or when the key is not found and the value of @p accum_or_assigns - * is `true`, nothing is changed and this operation is ignored. - * The algorithm assumes these situations occur while the key was modified or - * removed by other processes just now. - * - * @param n The number of key-value-score tuples to process. - * @param keys The keys to insert on GPU-accessible memory with shape (n). - * @param value_or_deltas The values or deltas to insert on GPU-accessible - * memory with shape (n, DIM). - * @param accum_or_assigns The operation type with shape (n). A value of - * `true` indicates to accum and `false` indicates to assign. - * @param scores The scores to insert on GPU-accessible memory with shape (n). - * @parblock - * The scores should be a `uint64_t` value for built-in strategies. For - * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. - * You can specify a value such as the timestamp of the key insertion or - * number of key occurrences to perform a customized eviction strategy. - * - * The @p scores should be `nullptr`, when the LRU eviction strategy is - * applied. - * @endparblock - * - * @param stream The CUDA stream that is used to execute the operation. - * - * @param ignore_evict_strategy A boolean option indicating whether if - * the accum_or_assign ignores the evict strategy of table with current - * scores anyway. If true, it does not check whether the scores confroms to - * the evict strategy. If false, it requires the scores follow the evict - * strategy of table. - */ - virtual void accum_or_assign(const size_type n, - const key_type* keys, // (n) - const value_type* value_or_deltas, // (n, DIM) - const bool* accum_or_assigns, // (n) - const score_type* scores = nullptr, // (n) - cudaStream_t stream = 0, - bool ignore_evict_strategy = false) = 0; - - /** - * @brief Searches the hash table for the specified keys. - * When a key is missing, the value in @p values and @p scores will be - * inserted. - * - * @param n The number of key-value-score tuples to search or insert. - * @param keys The keys to search on GPU-accessible memory with shape (n). - * @param values The values to search on GPU-accessible memory with - * shape (n, DIM). - * @param scores The scores to search on GPU-accessible memory with shape (n). - * @parblock - * If @p scores is `nullptr`, the score for each key will not be returned. - * @endparblock - * @param stream The CUDA stream that is used to execute the operation. - * @param unique_key If all keys in the same batch are unique. - * - */ - virtual void find_or_insert(const size_type n, const key_type* keys, // (n) - value_type* values, // (n * DIM) - score_type* scores = nullptr, // (n) - cudaStream_t stream = 0, bool unique_key = true, - bool ignore_evict_strategy = false) = 0; - - /** - * @brief Searches the hash table for the specified keys and returns address - * of the values. When a key is missing, the value in @p values and @p scores - * will be inserted. - * - * @warning This API returns internal addresses for high-performance but - * thread-unsafe. The caller is responsible for guaranteeing data consistency. - * - * @param n The number of key-value-score tuples to search or insert. - * @param keys The keys to search on GPU-accessible memory with shape (n). - * @param values The addresses of values to search on GPU-accessible memory - * with shape (n). - * @param founds The status that indicates if the keys are found on - * @param scores The scores to search on GPU-accessible memory with shape (n). - * @parblock - * If @p scores is `nullptr`, the score for each key will not be returned. - * @endparblock - * @param stream The CUDA stream that is used to execute the operation. - * @param unique_key If all keys in the same batch are unique. - * @param locked_key_ptrs If it isn't nullptr then the keys in the table will - * be locked, and key's address will write to locked_key_ptrs. Using - * unlock_keys to unlock these keys. - * - */ - virtual void find_or_insert(const size_type n, const key_type* keys, // (n) - value_type** values, // (n) - bool* founds, // (n) - score_type* scores = nullptr, // (n) - cudaStream_t stream = 0, bool unique_key = true, - bool ignore_evict_strategy = false, - key_type** locked_key_ptrs = nullptr) = 0; - - /** - * @brief - * This function will lock the keys in the table and unexisted keys will be - * ignored. - * - * @param n The number of keys in the table to be locked. - * @param locked_key_ptrs The pointers of locked keys in the table with shape - * (n). - * @param keys The keys to search on GPU-accessible memory with shape (n). - * @param succeededs The status that indicates if the lock operation is - * succeed. - * @param scores The scores of the input keys will set to scores if provided. - * @param stream The CUDA stream that is used to execute the operation. - * - */ - virtual void lock_keys(const size_type n, - key_type const* keys, // (n) - key_type** locked_key_ptrs, // (n) - bool* succeededs = nullptr, // (n) - cudaStream_t stream = 0, - score_type const* scores = nullptr) = 0; - - /** - * @brief Using pointers to address the keys in the hash table and set them - * to target keys. - * This function will unlock the keys in the table which are locked by - * the previous call to find_or_insert. - * - * @param n The number of keys in the table to be unlocked. - * @param locked_key_ptrs The pointers of locked keys in the table with shape - * (n). - * @param keys The keys to search on GPU-accessible memory with shape (n). - * @param succeededs The status that indicates if the unlock operation is - * succeed. - * @param stream The CUDA stream that is used to execute the operation. - * - */ - virtual void unlock_keys(const size_type n, - key_type** locked_key_ptrs, // (n) - const key_type* keys, // (n) - bool* succeededs = nullptr, // (n) - cudaStream_t stream = 0) = 0; - - /** - * @brief Assign new key-value-score tuples into the hash table. - * If the key doesn't exist, the operation on the key will be ignored. - * - * @param n Number of key-value-score tuples to insert or assign. - * @param keys The keys to insert on GPU-accessible memory with shape - * (n). - * @param values The values to insert on GPU-accessible memory with - * shape (n, DIM). - * @param scores The scores to insert on GPU-accessible memory with shape - * (n). - * @parblock - * The scores should be a `uint64_t` value for built-in strategies. For - * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. - * You can specify a value such as the timestamp of the key insertion or - * number of key occurrences to perform a customized eviction strategy. - * - * The @p scores should be `nullptr`, when the LRU eviction strategy is - * applied. - * @endparblock - * - * @param stream The CUDA stream that is used to execute the operation. - * - * @param unique_key If all keys in the same batch are unique. - */ - virtual void assign(const size_type n, - const key_type* keys, // (n) - const value_type* values, // (n, DIM) - const score_type* scores = nullptr, // (n) - cudaStream_t stream = 0, bool unique_key = true) = 0; - - /** - * @brief Assign new scores for keys. - * If the key doesn't exist, the operation on the key will be ignored. - * - * @param n Number of key-score pairs to assign. - * @param keys The keys to insert on GPU-accessible memory with shape - * (n). - * @parblock - * The scores should be a `uint64_t` value for built-in strategies. For - * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. - * You can specify a value such as the timestamp of the key insertion or - * number of key occurrences to perform a customized eviction strategy. - * - * The @p scores should be `nullptr`, when the LRU eviction strategy is - * applied. - * @endparblock - * - * @param stream The CUDA stream that is used to execute the operation. - * - * @param unique_key If all keys in the same batch are unique. - */ - virtual void assign_scores(const size_type n, - const key_type* keys, // (n) - const score_type* scores = nullptr, // (n) - cudaStream_t stream = 0, - bool unique_key = true) = 0; - - /** - * @brief Alias of `assign_scores`. - */ - virtual void assign(const size_type n, - const key_type* keys, // (n) - const score_type* scores = nullptr, // (n) - cudaStream_t stream = 0, bool unique_key = true) = 0; - - /** - * @brief Assign new values for each keys . - * If the key doesn't exist, the operation on the key will be ignored. - * - * @param n Number of key-value pairs to assign. - * @param keys The keys need to be operated, which must be on GPU-accessible - * memory with shape (n). - * @param values The values need to be updated, which must be on - * GPU-accessible memory with shape (n, DIM). - * - * @param stream The CUDA stream that is used to execute the operation. - * - * @param unique_key If all keys in the same batch are unique. - */ - virtual void assign_values(const size_type n, - const key_type* keys, // (n) - const value_type* values, // (n, DIM) - cudaStream_t stream = 0, - bool unique_key = true) = 0; - /** - * @brief Searches the hash table for the specified keys. - * - * @note When a key is missing, the value in @p values is not changed. - * - * @param n The number of key-value-score tuples to search. - * @param keys The keys to search on GPU-accessible memory with shape (n). - * @param values The values to search on GPU-accessible memory with - * shape (n, DIM). - * @param founds The status that indicates if the keys are found on - * GPU-accessible memory with shape (n). - * @param scores The scores to search on GPU-accessible memory with shape (n). - * @parblock - * If @p scores is `nullptr`, the score for each key will not be returned. - * @endparblock - * @param stream The CUDA stream that is used to execute the operation. - * - */ - virtual void find(const size_type n, const key_type* keys, // (n) - value_type* values, // (n, DIM) - bool* founds, // (n) - score_type* scores = nullptr, // (n) - cudaStream_t stream = 0) const = 0; - - /** - * @brief Searches the hash table for the specified keys. - * - * @note When the searched keys are not hit, missed keys/indices/size can be - * obtained. - * - * @param n The number of key-value-score tuples to search. - * @param keys The keys to search on GPU-accessible memory with shape (n). - * @param values The values to search on GPU-accessible memory with - * shape (n, DIM). - * @param missed_keys The missed keys to search on GPU-accessible memory with - * shape (n). - * @param missed_indices The missed indices to search on GPU-accessible memory - * with shape (n). - * @param missed_size The size of `missed_keys` and `missed_indices`. - * @param scores The scores to search on GPU-accessible memory with shape (n). - * @parblock - * If @p scores is `nullptr`, the score for each key will not be returned. - * @endparblock - * @param stream The CUDA stream that is used to execute the operation. - */ - virtual void find(const size_type n, const key_type* keys, // (n) - value_type* values, // (n, DIM) - key_type* missed_keys, // (n) - int* missed_indices, // (n) - int* missed_size, // scalar - score_type* scores = nullptr, // (n) - cudaStream_t stream = 0) const = 0; - - /** - * @brief Searches the hash table for the specified keys and returns address - * of the values. - * - * @note When a key is missing, the data in @p values won't change. - * @warning This API returns internal addresses for high-performance but - * thread-unsafe. The caller is responsible for guaranteeing data consistency. - * - * @param n The number of key-value-score tuples to search. - * @param keys The keys to search on GPU-accessible memory with shape (n). - * @param values The addresses of values to search on GPU-accessible memory - * with shape (n). - * @param founds The status that indicates if the keys are found on - * GPU-accessible memory with shape (n). - * @param scores The scores to search on GPU-accessible memory with shape (n). - * @parblock - * If @p scores is `nullptr`, the score for each key will not be returned. - * @endparblock - * @param stream The CUDA stream that is used to execute the operation. - * @param unique_key If all keys in the same batch are unique. - * - */ - virtual void find(const size_type n, const key_type* keys, // (n) - value_type** values, // (n) - bool* founds, // (n) - score_type* scores = nullptr, // (n) - cudaStream_t stream = 0, bool unique_key = true) const = 0; - - /** - * @brief Searches the hash table for the specified keys and returns address - * of the values, and will update the scores. - * - * @note When a key is missing, the data in @p values won't change. - * @warning This API returns internal addresses for high-performance but - * thread-unsafe. The caller is responsible for guaranteeing data consistency. - * - * @param n The number of key-value-score tuples to search. - * @param keys The keys to search on GPU-accessible memory with shape (n). - * @param values The addresses of values to search on GPU-accessible memory - * with shape (n). - * @param founds The status that indicates if the keys are found on - * GPU-accessible memory with shape (n). - * @param scores The scores to search on GPU-accessible memory with shape (n). - * @parblock - * If @p scores is `nullptr`, the score for each key will not be returned. - * @endparblock - * @param stream The CUDA stream that is used to execute the operation. - * @param unique_key If all keys in the same batch are unique. - * - */ - virtual void find_and_update(const size_type n, const key_type* keys, // (n) - value_type** values, // (n) - bool* founds, // (n) - score_type* scores = nullptr, // (n) - cudaStream_t stream = 0, - bool unique_key = true) = 0; - - /** - * @brief Checks if there are elements with key equivalent to `keys` in the - * table. - * - * @param n The number of `keys` to check. - * @param keys The keys to search on GPU-accessible memory with shape (n). - * @param founds The result that indicates if the keys are found, and should - * be allocated by caller on GPU-accessible memory with shape (n). - * @param stream The CUDA stream that is used to execute the operation. - * - */ - virtual void contains(const size_type n, const key_type* keys, // (n) - bool* founds, // (n) - cudaStream_t stream = 0) const = 0; - - /** - * @brief Removes specified elements from the hash table. - * - * @param n The number of keys to remove. - * @param keys The keys to remove on GPU-accessible memory. - * @param stream The CUDA stream that is used to execute the operation. - * - */ - virtual void erase(const size_type n, const key_type* keys, - cudaStream_t stream = 0) = 0; - - /** - * @brief Removes all of the elements in the hash table with no release - * object. - */ - virtual void clear(cudaStream_t stream = 0) = 0; - - /** - * @brief Exports a certain number of the key-value-score tuples from the - * hash table. - * - * @param n The maximum number of exported pairs. - * @param offset The position of the key to search. - * @param d_counter Accumulates amount of successfully exported values. - * @param keys The keys to dump from GPU-accessible memory with shape (n). - * @param values The values to dump from GPU-accessible memory with shape - * (n, DIM). - * @param scores The scores to search on GPU-accessible memory with shape (n). - * @parblock - * If @p scores is `nullptr`, the score for each key will not be returned. - * @endparblock - * - * @param stream The CUDA stream that is used to execute the operation. - * - * @return The number of elements dumped. - * - * @throw CudaException If the key-value size is too large for GPU shared - * memory. Reducing the value for @p n is currently required if this exception - * occurs. - */ - virtual void export_batch(size_type n, const size_type offset, - size_type* d_counter, // (1) - key_type* keys, // (n) - value_type* values, // (n, DIM) - score_type* scores = nullptr, // (n) - cudaStream_t stream = 0) const = 0; - - virtual size_type export_batch(const size_type n, const size_type offset, - key_type* keys, // (n) - value_type* values, // (n, DIM) - score_type* scores = nullptr, // (n) - cudaStream_t stream = 0) const = 0; - - /** - * @brief Indicates if the hash table has no elements. - * - * @param stream The CUDA stream that is used to execute the operation. - * @return `true` if the table is empty and `false` otherwise. - */ - virtual bool empty(cudaStream_t stream = 0) const = 0; - - /** - * @brief Returns the hash table size. - * - * @param stream The CUDA stream that is used to execute the operation. - * @return The table size. - */ - virtual size_type size(cudaStream_t stream = 0) const = 0; - - /** - * @brief Returns the hash table capacity. - * - * @note The value that is returned might be less than the actual capacity of - * the hash table because the hash table currently keeps the capacity to be - * a power of 2 for performance considerations. - * - * @return The table capacity. - */ - virtual size_type capacity() const = 0; - - /** - * @brief Sets the number of buckets to the number that is needed to - * accommodate at least @p new_capacity elements without exceeding the maximum - * load factor. This method rehashes the hash table. Rehashing puts the - * elements into the appropriate buckets considering that total number of - * buckets has changed. - * - * @note If the value of @p new_capacity or double of @p new_capacity is - * greater or equal than `options_.max_capacity`, the reserve does not perform - * any change to the hash table. - * - * @param new_capacity The requested capacity for the hash table. - * @param stream The CUDA stream that is used to execute the operation. - */ - virtual void reserve(const size_type new_capacity, - cudaStream_t stream = 0) = 0; - - /** - * @brief Returns the average number of elements per slot, that is, size() - * divided by capacity(). - * - * @param stream The CUDA stream that is used to execute the operation. - * - * @return The load factor - */ - virtual float load_factor(cudaStream_t stream = 0) const = 0; - - /** - * @brief Set max_capacity of the table. - * - * @param new_max_capacity The new expecting max_capacity. It must be power - * of 2. Otherwise it will raise an error. - */ - virtual void set_max_capacity(size_type new_max_capacity) = 0; - - /** - * @brief Returns the dimension of the vectors. - * - * @return The dimension of the vectors. - */ - virtual size_type dim() const noexcept = 0; - - /** - * @brief Returns The length of each bucket. - * - * @return The length of each bucket. - */ - virtual size_type max_bucket_size() const noexcept = 0; - - /** - * @brief Returns the number of buckets in the table. - * - * @return The number of buckets in the table. - */ - virtual size_type bucket_count() const noexcept = 0; - - /** - * @brief Save keys, vectors, scores in table to file or files. - * - * @param file A BaseKVFile object defined the file format on host filesystem. - * @param max_workspace_size Saving is conducted in chunks. This value denotes - * the maximum amount of temporary memory to use when dumping the table. - * Larger values *can* lead to higher performance. - * @param stream The CUDA stream used to execute the operation. - * - * @return Number of KV pairs saved to file. - */ - virtual size_type save(BaseKVFile* file, - const size_t max_workspace_size = 1L * 1024 * 1024, - cudaStream_t stream = 0) const = 0; - - /** - * @brief Load keys, vectors, scores from file to table. - * - * @param file An BaseKVFile defined the file format within filesystem. - * @param max_workspace_size Loading is conducted in chunks. This value - * denotes the maximum size of such chunks. Larger values *can* lead to higher - * performance. - * @param stream The CUDA stream used to execute the operation. - * - * @return Number of keys loaded from file. - */ - virtual size_type load(BaseKVFile* file, - const size_t max_workspace_size = 1L * 1024 * 1024, - cudaStream_t stream = 0) = 0; - - virtual void set_global_epoch(const uint64_t epoch) = 0; -}; - /** * A HierarchicalKV hash table is a concurrent and hierarchical hash table that * is powered by GPUs and can use HBM and host memory as storage for key-value @@ -906,6 +207,10 @@ class HashTableBase { * Supported types: `uint64_t` and `uint32_t` (only for * `EvictStrategy::kCustomized`). * + * @note ArchTag controls internal tuning (SM-specific pipeline config), not + * the actual compiled GPU binary architecture. We keep the default at + * Sm80 so that all kernels reuse the existing specializations, while + * nvcc still generates sm_100 code via -gencode flags. */ template @@ -1240,7 +545,7 @@ class HashTable : public HashTableBase { constexpr uint32_t BLOCK_SIZE = 128; upsert_kernel_lock_key_hybrid + BLOCK_SIZE, evict_strategy> <<<(n + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>( table_->buckets, table_->buckets_size, table_->buckets_num, options_.max_bucket_size, options_.dim, keys, d_dst, scores, @@ -1251,7 +556,7 @@ class HashTable : public HashTableBase { const size_t N = n * TILE_SIZE; const size_t grid_size = SAFE_GET_GRID_SIZE(N, block_size); - upsert_kernel<<>>( d_table_, table_->buckets, options_.max_bucket_size, table_->buckets_num, options_.dim, keys, d_dst, scores, @@ -3871,6 +3176,18 @@ class HashTable : public HashTableBase { cudaMemcpy(d_table_, table_, sizeof(TableCore), cudaMemcpyDefault)); } + public: + // Expose device buckets and layout for read-only lookup kernels + inline nv::merlin::Bucket* device_buckets() const { + return table_ ? table_->buckets : nullptr; + } + inline size_t device_bucket_count() const { + return table_ ? table_->buckets_num : 0; + } + inline size_t device_bucket_max_size() const { + return table_ ? table_->bucket_max_size : 0; + } + private: HashTableOptions options_; TableCore* table_ = nullptr; diff --git a/include/merlin_hashtable_base.hpp b/include/merlin_hashtable_base.hpp new file mode 100644 index 00000000..b9c25ebb --- /dev/null +++ b/include/merlin_hashtable_base.hpp @@ -0,0 +1,737 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +namespace nv { +namespace merlin { + +struct HashTableOptions; +class BaseAllocator; + +template +struct Bucket; + +template +class BaseKVFile; + +template +class HashTableBase { + public: + using size_type = size_t; + using key_type = K; + using value_type = V; + using score_type = S; + using allocator_type = BaseAllocator; + using bucket_type = nv::merlin::Bucket; + + public: + virtual ~HashTableBase() {} + + /** + * @brief Initialize a merlin::HashTable. + * + * @param options The configuration options. + */ + virtual void init(const HashTableOptions& options, + allocator_type* allocator = nullptr) = 0; + + /** + * @brief Insert new key-value-score tuples into the hash table. + * If the key already exists, the values and scores are assigned new values. + * + * If the target bucket is full, the keys with minimum score will be + * overwritten by new key unless the score of the new key is even less than + * minimum score of the target bucket. + * + * @param n Number of key-value-score tuples to insert or assign. + * @param keys The keys to insert on GPU-accessible memory with shape + * (n). + * @param values The values to insert on GPU-accessible memory with + * shape (n, DIM). + * @param scores The scores to insert on GPU-accessible memory with shape + * (n). + * @parblock + * The scores should be a `uint64_t` value for built-in strategies. For + * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. + * You can specify a value such as the timestamp of the key insertion or + * number of key occurrences to perform a customized eviction strategy. + * + * The @p scores should be `nullptr`, when the LRU eviction strategy is + * applied. + * @endparblock + * + * @param stream The CUDA stream that is used to execute the operation. + * @param unique_key If all keys in the same batch are unique. + * + * @param ignore_evict_strategy A boolean option indicating whether if + * the insert_or_assign ignores the evict strategy of table with current + * scores anyway. If true, it does not check whether the scores conforms to + * the evict strategy. If false, it requires the scores follow the evict + * strategy of table. + */ + virtual void insert_or_assign(const size_type n, + const key_type* keys, // (n) + const value_type* values, // (n, DIM) + const score_type* scores = nullptr, // (n) + cudaStream_t stream = 0, bool unique_key = true, + bool ignore_evict_strategy = false) = 0; + + /** + * @brief Insert new key-value-score tuples into the hash table. + * If the key already exists, the values and scores are assigned new values. + * + * If the target bucket is full, the keys with minimum score will be + * overwritten by new key unless the score of the new key is even less than + * minimum score of the target bucket. The overwritten key with minimum + * score will be evicted, with its values and score, to evicted_keys, + * evicted_values, evcted_scores seperately in compact format. + * + * @param n Number of key-value-score tuples to insert or assign. + * @param keys The keys to insert on GPU-accessible memory with shape + * (n). + * @param values The values to insert on GPU-accessible memory with + * shape (n, DIM). + * @param scores The scores to insert on GPU-accessible memory with shape + * (n). + * @param scores The scores to insert on GPU-accessible memory with shape + * (n). + * @params evicted_keys The output of keys replaced with minimum score. + * @params evicted_values The output of values replaced with minimum score on + * keys. + * @params evicted_scores The output of scores replaced with minimum score on + * keys. + * @parblock + * The scores should be a `uint64_t` value for built-in strategies. For + * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. + * You can specify a value such as the timestamp of the key insertion or + * number of key occurrences to perform a customized eviction strategy. + * + * The @p scores should be `nullptr`, when the LRU eviction strategy is + * applied. + * @endparblock + * + * @param d_evicted_counter The number of elements evicted on GPU-accessible + * memory. @notice The caller should guarantee it is set to `0` before + * calling. + * @param stream The CUDA stream that is used to execute the operation. + * @param unique_key If all keys in the same batch are unique. + * + * @param ignore_evict_strategy A boolean option indicating whether if + * the insert_or_assign ignores the evict strategy of table with current + * scores anyway. If true, it does not check whether the scores confroms to + * the evict strategy. If false, it requires the scores follow the evict + * strategy of table. + */ + virtual void insert_and_evict(const size_type n, + const key_type* keys, // (n) + const value_type* values, // (n, DIM) + const score_type* scores, // (n) + key_type* evicted_keys, // (n) + value_type* evicted_values, // (n, DIM) + score_type* evicted_scores, // (n) + size_type* d_evicted_counter, // (1) + cudaStream_t stream = 0, bool unique_key = true, + bool ignore_evict_strategy = false) = 0; + + /** + * @brief Insert new key-value-score tuples into the hash table. + * If the key already exists, the values and scores are assigned new values. + * + * If the target bucket is full, the keys with minimum score will be + * overwritten by new key unless the score of the new key is even less than + * minimum score of the target bucket. The overwritten key with minimum + * score will be evicted, with its values and score, to evicted_keys, + * evicted_values, evcted_scores seperately in compact format. + * + * @param n Number of key-value-score tuples to insert or assign. + * @param keys The keys to insert on GPU-accessible memory with shape + * (n). + * @param values The values to insert on GPU-accessible memory with + * shape (n, DIM). + * @param scores The scores to insert on GPU-accessible memory with shape + * (n). + * @param scores The scores to insert on GPU-accessible memory with shape + * (n). + * @params evicted_keys The output of keys replaced with minimum score. + * @params evicted_values The output of values replaced with minimum score on + * keys. + * @params evicted_scores The output of scores replaced with minimum score on + * keys. + * @parblock + * The scores should be a `uint64_t` value for built-in strategies. For + * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. + * You can specify a value such as the timestamp of the key insertion or + * number of key occurrences to perform a customized eviction strategy. + * + * The @p scores should be `nullptr`, when the LRU eviction strategy is + * applied. + * @endparblock + * + * @param stream The CUDA stream that is used to execute the operation. + * @param unique_key If all keys in the same batch are unique. + * + * @param ignore_evict_strategy A boolean option indicating whether if + * the insert_or_assign ignores the evict strategy of table with current + * scores anyway. If true, it does not check whether the scores confroms to + * the evict strategy. If false, it requires the scores follow the evict + * strategy of table. + * + * @return The number of elements evicted. + */ + virtual size_type insert_and_evict(const size_type n, + const key_type* keys, // (n) + const value_type* values, // (n, DIM) + const score_type* scores, // (n) + key_type* evicted_keys, // (n) + value_type* evicted_values, // (n, DIM) + score_type* evicted_scores, // (n) + cudaStream_t stream = 0, + bool unique_key = true, + bool ignore_evict_strategy = false) = 0; + + /** + * Searches for each key in @p keys in the hash table. + * If the key is found and the corresponding value in @p accum_or_assigns is + * `true`, the @p vectors_or_deltas is treated as a delta to the old + * value, and the delta is added to the old value of the key. + * + * If the key is not found and the corresponding value in @p accum_or_assigns + * is `false`, the @p vectors_or_deltas is treated as a new value and the + * key-value pair is updated in the table directly. + * + * @note When the key is found and the value of @p accum_or_assigns is + * `false`, or when the key is not found and the value of @p accum_or_assigns + * is `true`, nothing is changed and this operation is ignored. + * The algorithm assumes these situations occur while the key was modified or + * removed by other processes just now. + * + * @param n The number of key-value-score tuples to process. + * @param keys The keys to insert on GPU-accessible memory with shape (n). + * @param value_or_deltas The values or deltas to insert on GPU-accessible + * memory with shape (n, DIM). + * @param accum_or_assigns The operation type with shape (n). A value of + * `true` indicates to accum and `false` indicates to assign. + * @param scores The scores to insert on GPU-accessible memory with shape (n). + * @parblock + * The scores should be a `uint64_t` value for built-in strategies. For + * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. + * You can specify a value such as the timestamp of the key insertion or + * number of key occurrences to perform a customized eviction strategy. + * + * The @p scores should be `nullptr`, when the LRU eviction strategy is + * applied. + * @endparblock + * + * @param stream The CUDA stream that is used to execute the operation. + * + * @param ignore_evict_strategy A boolean option indicating whether if + * the accum_or_assign ignores the evict strategy of table with current + * scores anyway. If true, it does not check whether the scores confroms to + * the evict strategy. If false, it requires the scores follow the evict + * strategy of table. + */ + virtual void accum_or_assign(const size_type n, + const key_type* keys, // (n) + const value_type* value_or_deltas, // (n, DIM) + const bool* accum_or_assigns, // (n) + const score_type* scores = nullptr, // (n) + cudaStream_t stream = 0, + bool ignore_evict_strategy = false) = 0; + + /** + * @brief Searches the hash table for the specified keys. + * When a key is missing, the value in @p values and @p scores will be + * inserted. + * + * @param n The number of key-value-score tuples to search or insert. + * @param keys The keys to search on GPU-accessible memory with shape (n). + * @param values The values to search on GPU-accessible memory with + * shape (n, DIM). + * @param scores The scores to search on GPU-accessible memory with shape (n). + * @parblock + * If @p scores is `nullptr`, the score for each key will not be returned. + * @endparblock + * @param stream The CUDA stream that is used to execute the operation. + * @param unique_key If all keys in the same batch are unique. + * + */ + virtual void find_or_insert(const size_type n, const key_type* keys, // (n) + value_type* values, // (n * DIM) + score_type* scores = nullptr, // (n) + cudaStream_t stream = 0, bool unique_key = true, + bool ignore_evict_strategy = false) = 0; + + /** + * @brief Searches the hash table for the specified keys and returns address + * of the values. When a key is missing, the value in @p values and @p scores + * will be inserted. + * + * @warning This API returns internal addresses for high-performance but + * thread-unsafe. The caller is responsible for guaranteeing data consistency. + * + * @param n The number of key-value-score tuples to search or insert. + * @param keys The keys to search on GPU-accessible memory with shape (n). + * @param values The addresses of values to search on GPU-accessible memory + * with shape (n). + * @param founds The status that indicates if the keys are found on + * @param scores The scores to search on GPU-accessible memory with shape (n). + * @parblock + * If @p scores is `nullptr`, the score for each key will not be returned. + * @endparblock + * @param stream The CUDA stream that is used to execute the operation. + * @param unique_key If all keys in the same batch are unique. + * @param locked_key_ptrs If it isn't nullptr then the keys in the table will + * be locked, and key's address will write to locked_key_ptrs. Using + * unlock_keys to unlock these keys. + * + */ + virtual void find_or_insert(const size_type n, const key_type* keys, // (n) + value_type** values, // (n) + bool* founds, // (n) + score_type* scores = nullptr, // (n) + cudaStream_t stream = 0, bool unique_key = true, + bool ignore_evict_strategy = false, + key_type** locked_key_ptrs = nullptr) = 0; + + /** + * @brief + * This function will lock the keys in the table and unexisted keys will be + * ignored. + * + * @param n The number of keys in the table to be locked. + * @param locked_key_ptrs The pointers of locked keys in the table with shape + * (n). + * @param keys The keys to search on GPU-accessible memory with shape (n). + * @param succeededs The status that indicates if the lock operation is + * succeed. + * @param scores The scores of the input keys will set to scores if provided. + * @param stream The CUDA stream that is used to execute the operation. + * + */ + virtual void lock_keys(const size_type n, + key_type const* keys, // (n) + key_type** locked_key_ptrs, // (n) + bool* succeededs = nullptr, // (n) + cudaStream_t stream = 0, + score_type const* scores = nullptr) = 0; + + /** + * @brief Using pointers to address the keys in the hash table and set them + * to target keys. + * This function will unlock the keys in the table which are locked by + * the previous call to find_or_insert. + * + * @param n The number of keys in the table to be unlocked. + * @param locked_key_ptrs The pointers of locked keys in the table with shape + * (n). + * @param keys The keys to search on GPU-accessible memory with shape (n). + * @param succeededs The status that indicates if the unlock operation is + * succeed. + * @param stream The CUDA stream that is used to execute the operation. + * + */ + virtual void unlock_keys(const size_type n, + key_type** locked_key_ptrs, // (n) + const key_type* keys, // (n) + bool* succeededs = nullptr, // (n) + cudaStream_t stream = 0) = 0; + + /** + * @brief Assign new key-value-score tuples into the hash table. + * If the key doesn't exist, the operation on the key will be ignored. + * + * @param n Number of key-value-score tuples to insert or assign. + * @param keys The keys to insert on GPU-accessible memory with shape + * (n). + * @param values The values to insert on GPU-accessible memory with + * shape (n, DIM). + * @param scores The scores to insert on GPU-accessible memory with shape + * (n). + * @parblock + * The scores should be a `uint64_t` value for built-in strategies. For + * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. + * You can specify a value such as the timestamp of the key insertion or + * number of key occurrences to perform a customized eviction strategy. + * + * The @p scores should be `nullptr`, when the LRU eviction strategy is + * applied. + * @endparblock + * + * @param stream The CUDA stream that is used to execute the operation. + * + * @param unique_key If all keys in the same batch are unique. + */ + virtual void assign(const size_type n, + const key_type* keys, // (n) + const value_type* values, // (n, DIM) + const score_type* scores = nullptr, // (n) + cudaStream_t stream = 0, bool unique_key = true) = 0; + + /** + * @brief Assign new scores for keys. + * If the key doesn't exist, the operation on the key will be ignored. + * + * @param n Number of key-score pairs to assign. + * @param keys The keys to insert on GPU-accessible memory with shape + * (n). + * @parblock + * The scores should be a `uint64_t` value for built-in strategies. For + * `EvictStrategy::kCustomized`, `uint32_t` scores are also supported. + * You can specify a value such as the timestamp of the key insertion or + * number of key occurrences to perform a customized eviction strategy. + * + * The @p scores should be `nullptr`, when the LRU eviction strategy is + * applied. + * @endparblock + * + * @param stream The CUDA stream that is used to execute the operation. + * + * @param unique_key If all keys in the same batch are unique. + */ + virtual void assign_scores(const size_type n, + const key_type* keys, // (n) + const score_type* scores = nullptr, // (n) + cudaStream_t stream = 0, + bool unique_key = true) = 0; + + /** + * @brief Alias of `assign_scores`. + */ + virtual void assign(const size_type n, + const key_type* keys, // (n) + const score_type* scores = nullptr, // (n) + cudaStream_t stream = 0, bool unique_key = true) = 0; + + /** + * @brief Assign new values for each keys . + * If the key doesn't exist, the operation on the key will be ignored. + * + * @param n Number of key-value pairs to assign. + * @param keys The keys need to be operated, which must be on GPU-accessible + * memory with shape (n). + * @param values The values need to be updated, which must be on + * GPU-accessible memory with shape (n, DIM). + * + * @param stream The CUDA stream that is used to execute the operation. + * + * @param unique_key If all keys in the same batch are unique. + */ + virtual void assign_values(const size_type n, + const key_type* keys, // (n) + const value_type* values, // (n, DIM) + cudaStream_t stream = 0, + bool unique_key = true) = 0; + /** + * @brief Searches the hash table for the specified keys. + * + * @note When a key is missing, the value in @p values is not changed. + * + * @param n The number of key-value-score tuples to search. + * @param keys The keys to search on GPU-accessible memory with shape (n). + * @param values The values to search on GPU-accessible memory with + * shape (n, DIM). + * @param founds The status that indicates if the keys are found on + * GPU-accessible memory with shape (n). + * @param scores The scores to search on GPU-accessible memory with shape (n). + * @parblock + * If @p scores is `nullptr`, the score for each key will not be returned. + * @endparblock + * @param stream The CUDA stream that is used to execute the operation. + * + */ + virtual void find(const size_type n, const key_type* keys, // (n) + value_type* values, // (n, DIM) + bool* founds, // (n) + score_type* scores = nullptr, // (n) + cudaStream_t stream = 0) const = 0; + + /** + * @brief Searches the hash table for the specified keys. + * + * @note When the searched keys are not hit, missed keys/indices/size can be + * obtained. + * + * @param n The number of key-value-score tuples to search. + * @param keys The keys to search on GPU-accessible memory with shape (n). + * @param values The values to search on GPU-accessible memory with + * shape (n, DIM). + * @param missed_keys The missed keys to search on GPU-accessible memory with + * shape (n). + * @param missed_indices The missed indices to search on GPU-accessible memory + * with shape (n). + * @param missed_size The size of `missed_keys` and `missed_indices`. + * @param scores The scores to search on GPU-accessible memory with shape (n). + * @parblock + * If @p scores is `nullptr`, the score for each key will not be returned. + * @endparblock + * @param stream The CUDA stream that is used to execute the operation. + */ + virtual void find(const size_type n, const key_type* keys, // (n) + value_type* values, // (n, DIM) + key_type* missed_keys, // (n) + int* missed_indices, // (n) + int* missed_size, // scalar + score_type* scores = nullptr, // (n) + cudaStream_t stream = 0) const = 0; + + /** + * @brief Searches the hash table for the specified keys and returns address + * of the values. + * + * @note When a key is missing, the data in @p values won't change. + * @warning This API returns internal addresses for high-performance but + * thread-unsafe. The caller is responsible for guaranteeing data consistency. + * + * @param n The number of key-value-score tuples to search. + * @param keys The keys to search on GPU-accessible memory with shape (n). + * @param values The addresses of values to search on GPU-accessible memory + * with shape (n). + * @param founds The status that indicates if the keys are found on + * GPU-accessible memory with shape (n). + * @param scores The scores to search on GPU-accessible memory with shape (n). + * @parblock + * If @p scores is `nullptr`, the score for each key will not be returned. + * @endparblock + * @param stream The CUDA stream that is used to execute the operation. + * @param unique_key If all keys in the same batch are unique. + * + */ + virtual void find(const size_type n, const key_type* keys, // (n) + value_type** values, // (n) + bool* founds, // (n) + score_type* scores = nullptr, // (n) + cudaStream_t stream = 0, bool unique_key = true) const = 0; + + /** + * @brief Searches the hash table for the specified keys and returns address + * of the values, and will update the scores. + * + * @note When a key is missing, the data in @p values won't change. + * @warning This API returns internal addresses for high-performance but + * thread-unsafe. The caller is responsible for guaranteeing data consistency. + * + * @param n The number of key-value-score tuples to search. + * @param keys The keys to search on GPU-accessible memory with shape (n). + * @param values The addresses of values to search on GPU-accessible memory + * with shape (n). + * @param founds The status that indicates if the keys are found on + * GPU-accessible memory with shape (n). + * @param scores The scores to search on GPU-accessible memory with shape (n). + * @parblock + * If @p scores is `nullptr`, the score for each key will not be returned. + * @endparblock + * @param stream The CUDA stream that is used to execute the operation. + * @param unique_key If all keys in the same batch are unique. + * + */ + virtual void find_and_update(const size_type n, const key_type* keys, // (n) + value_type** values, // (n) + bool* founds, // (n) + score_type* scores = nullptr, // (n) + cudaStream_t stream = 0, + bool unique_key = true) = 0; + + /** + * @brief Checks if there are elements with key equivalent to `keys` in the + * table. + * + * @param n The number of `keys` to check. + * @param keys The keys to search on GPU-accessible memory with shape (n). + * @param founds The result that indicates if the keys are found, and should + * be allocated by caller on GPU-accessible memory with shape (n). + * @param stream The CUDA stream that is used to execute the operation. + * + */ + virtual void contains(const size_type n, const key_type* keys, // (n) + bool* founds, // (n) + cudaStream_t stream = 0) const = 0; + + /** + * @brief Removes specified elements from the hash table. + * + * @param n The number of keys to remove. + * @param keys The keys to remove on GPU-accessible memory. + * @param stream The CUDA stream that is used to execute the operation. + * + */ + virtual void erase(const size_type n, const key_type* keys, + cudaStream_t stream = 0) = 0; + + /** + * @brief Removes all of the elements in the hash table with no release + * object. + */ + virtual void clear(cudaStream_t stream = 0) = 0; + + /** + * @brief Exports a certain number of the key-value-score tuples from the + * hash table. + * + * @param n The maximum number of exported pairs. + * @param offset The position of the key to search. + * @param d_counter Accumulates amount of successfully exported values. + * @param keys The keys to dump from GPU-accessible memory with shape (n). + * @param values The values to dump from GPU-accessible memory with shape + * (n, DIM). + * @param scores The scores to search on GPU-accessible memory with shape (n). + * @parblock + * If @p scores is `nullptr`, the score for each key will not be returned. + * @endparblock + * + * @param stream The CUDA stream that is used to execute the operation. + * + * @return The number of elements dumped. + * + * @throw CudaException If the key-value size is too large for GPU shared + * memory. Reducing the value for @p n is currently required if this exception + * occurs. + */ + virtual void export_batch(size_type n, const size_type offset, + size_type* d_counter, // (1) + key_type* keys, // (n) + value_type* values, // (n, DIM) + score_type* scores = nullptr, // (n) + cudaStream_t stream = 0) const = 0; + + virtual size_type export_batch(const size_type n, const size_type offset, + key_type* keys, // (n) + value_type* values, // (n, DIM) + score_type* scores = nullptr, // (n) + cudaStream_t stream = 0) const = 0; + + /** + * @brief Indicates if the hash table has no elements. + * + * @param stream The CUDA stream that is used to execute the operation. + * @return `true` if the table is empty and `false` otherwise. + */ + virtual bool empty(cudaStream_t stream = 0) const = 0; + + /** + * @brief Returns the hash table size. + * + * @param stream The CUDA stream that is used to execute the operation. + * @return The table size. + */ + virtual size_type size(cudaStream_t stream = 0) const = 0; + + /** + * @brief Returns the hash table capacity. + * + * @note The value that is returned might be less than the actual capacity of + * the hash table because the hash table currently keeps the capacity to be + * a power of 2 for performance considerations. + * + * @return The table capacity. + */ + virtual size_type capacity() const = 0; + + /** + * @brief Sets the number of buckets to the number that is needed to + * accommodate at least @p new_capacity elements without exceeding the maximum + * load factor. This method rehashes the hash table. Rehashing puts the + * elements into the appropriate buckets considering that total number of + * buckets has changed. + * + * @note If the value of @p new_capacity or double of @p new_capacity is + * greater or equal than `options_.max_capacity`, the reserve does not perform + * any change to the hash table. + * + * @param new_capacity The requested capacity for the hash table. + * @param stream The CUDA stream that is used to execute the operation. + */ + virtual void reserve(const size_type new_capacity, + cudaStream_t stream = 0) = 0; + + /** + * @brief Returns the average number of elements per slot, that is, size() + * divided by capacity(). + * + * @param stream The CUDA stream that is used to execute the operation. + * + * @return The load factor + */ + virtual float load_factor(cudaStream_t stream = 0) const = 0; + + /** + * @brief Set max_capacity of the table. + * + * @param new_max_capacity The new expecting max_capacity. It must be power + * of 2. Otherwise it will raise an error. + */ + virtual void set_max_capacity(size_type new_max_capacity) = 0; + + /** + * @brief Returns the dimension of the vectors. + * + * @return The dimension of the vectors. + */ + virtual size_type dim() const noexcept = 0; + + /** + * @brief Returns The length of each bucket. + * + * @return The length of each bucket. + */ + virtual size_type max_bucket_size() const noexcept = 0; + + /** + * @brief Returns the number of buckets in the table. + * + * @return The number of buckets in the table. + */ + virtual size_type bucket_count() const noexcept = 0; + + /** + * @brief Save keys, vectors, scores in table to file or files. + * + * @param file A BaseKVFile object defined the file format on host filesystem. + * @param max_workspace_size Saving is conducted in chunks. This value denotes + * the maximum amount of temporary memory to use when dumping the table. + * Larger values *can* lead to higher performance. + * @param stream The CUDA stream used to execute the operation. + * + * @return Number of KV pairs saved to file. + */ + virtual size_type save(BaseKVFile* file, + const size_t max_workspace_size = 1L * 1024 * 1024, + cudaStream_t stream = 0) const = 0; + + /** + * @brief Load keys, vectors, scores from file to table. + * + * @param file An BaseKVFile defined the file format within filesystem. + * @param max_workspace_size Loading is conducted in chunks. This value + * denotes the maximum size of such chunks. Larger values *can* lead to higher + * performance. + * @param stream The CUDA stream used to execute the operation. + * + * @return Number of keys loaded from file. + */ + virtual size_type load(BaseKVFile* file, + const size_t max_workspace_size = 1L * 1024 * 1024, + cudaStream_t stream = 0) = 0; + + virtual void set_global_epoch(const uint64_t epoch) = 0; +}; + +} // namespace merlin +} // namespace nv diff --git a/include/merlin_hashtable_device.cuh b/include/merlin_hashtable_device.cuh new file mode 100644 index 00000000..41a1b7f3 --- /dev/null +++ b/include/merlin_hashtable_device.cuh @@ -0,0 +1,158 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include "merlin/core_kernels/kernel_utils.cuh" +#include "merlin_hashtable.cuh" +#include "merlin_hashtable_base.hpp" + +namespace nv { +namespace merlin { +namespace device { + +/** + * @brief Lightweight device-side view of a hash table. + * + * This view exposes only the metadata and bucket pointer needed by device + * read-only APIs. + */ +template +struct HashTableDeviceView { + using base_type = HashTableBase; + using key_type = typename base_type::key_type; + using value_type = typename base_type::value_type; + using score_type = typename base_type::score_type; + using bucket_type = typename base_type::bucket_type; + + // Read-only view. Non-const to match find_without_lock_readonly_no_sync. + bucket_type* buckets; + size_t bucket_count; + size_t bucket_max_size; +}; + +/** + * @brief Compute bucket index and aligned start for a key. + * + * @tparam K Key type. + * @tparam V Value type. + * @tparam S Score type. + * @tparam TILE_SIZE Alignment granularity for the start position. + * @param view Device-side view of the table. + * @param key Key to hash and locate. + * @param bucket_idx Output bucket index. + * @param aligned_start Output aligned start position in the bucket. + * @return Always returns true. The caller must ensure view metadata is valid. + */ +template +__device__ __forceinline__ bool compute_bucket_index_and_aligned_start( + const HashTableDeviceView& view, + const typename HashTableDeviceView::key_type key, + uint32_t* bucket_idx, uint32_t* aligned_start) { + static_assert((TILE_SIZE & (TILE_SIZE - 1)) == 0, + "TILE_SIZE must be power of two."); + // bucket_max_size is guaranteed to be power-of-two. + const uint32_t bucket_capacity = static_cast(view.bucket_max_size); + const uint64_t total_slots = + static_cast(view.bucket_count) * bucket_capacity; + const uint64_t hashed_key = static_cast(Murmur3HashDevice(key)); + const uint64_t global_idx = hashed_key % total_slots; + const uint32_t bucket_shift = + static_cast(__ffs(bucket_capacity) - 1); + *bucket_idx = static_cast(global_idx >> bucket_shift); + uint32_t start_idx = + static_cast(global_idx) & (bucket_capacity - 1); + *aligned_start = align_to(start_idx); + return true; +} + +/** + * @brief Read-only lookup without tile synchronization in a known bucket. + * + * @tparam K Key type. + * @tparam V Value type. + * @tparam S Score type. + * @tparam TILE_SIZE Tile size used for probing. + * @param bucket Bucket pointer for the target bucket. + * @param key Key to search in the bucket. + * @param aligned_start Aligned start position in the bucket. + * @param bucket_max_size Bucket capacity (must be power-of-two). + * @param rank Thread rank inside the tile. + * @note The table must be immutable during lookup, and keys in the same batch + * must be non-overlapping. + * @return Position inside the bucket if found, otherwise `-1`. + */ +template +__device__ __forceinline__ int find_readonly_no_sync_in_bucket( + typename HashTableDeviceView::bucket_type* bucket, + const typename HashTableDeviceView::key_type key, + const uint32_t aligned_start, const uint32_t bucket_max_size, + const uint32_t rank) { + return find_without_lock_readonly_no_sync( + bucket, key, aligned_start, bucket_max_size, rank); +} + +/** + * @brief Read-only lookup without tile synchronization. + * + * @tparam K Key type. + * @tparam V Value type. + * @tparam S Score type. + * @tparam TILE_SIZE Tile size used for probing. + * @param view Device-side view of the table. + * @param key Key to search in the table. + * @param rank Thread rank inside the tile. + * @note The table must be immutable during lookup, and keys in the same batch + * must be non-overlapping. + * @return Position inside the bucket if found, otherwise `-1`. + */ +template +__device__ __forceinline__ int find_readonly_no_sync( + const HashTableDeviceView& view, + const typename HashTableDeviceView::key_type key, + const uint32_t rank) { + uint32_t bucket_idx = 0; + uint32_t aligned_start = 0; + if (!compute_bucket_index_and_aligned_start( + view, key, &bucket_idx, &aligned_start)) { + return -1; + } + return find_readonly_no_sync_in_bucket( + view.buckets + bucket_idx, key, aligned_start, + static_cast(view.bucket_max_size), rank); +} + +/** + * @brief Create a device-side view from a host table instance. + * + * @param table Host-side hash table. + * @return Device-side view that can be used in kernels. + */ +template +__host__ inline HashTableDeviceView make_device_view( + const HashTable& table) { + HashTableDeviceView view{}; + view.buckets = table.device_buckets(); + view.bucket_count = table.device_bucket_count(); + view.bucket_max_size = table.device_bucket_max_size(); + return view; +} + +} // namespace device +} // namespace merlin +} // namespace nv diff --git a/tests/merlin_hashtable_device_api_test.cc.cu b/tests/merlin_hashtable_device_api_test.cc.cu new file mode 100644 index 00000000..4c3cbd38 --- /dev/null +++ b/tests/merlin_hashtable_device_api_test.cc.cu @@ -0,0 +1,221 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include "merlin_hashtable.cuh" +#include "merlin_hashtable_device.cuh" +#include "test_util.cuh" + +namespace cg = cooperative_groups; + +constexpr uint32_t kTileSize = 4; +constexpr size_t kDim = 4; +using K = uint64_t; +using V = float; +using S = uint64_t; +using TableOptions = nv::merlin::HashTableOptions; +using EvictStrategy = nv::merlin::EvictStrategy; +using Table = nv::merlin::HashTable; + +template +__global__ void ReadonlyLookupKernel( + nv::merlin::device::HashTableDeviceView view, const K* keys, + int key_count, bool* founds, int* positions, bool* key_matches) { + cg::thread_block_tile tile = + cg::tiled_partition(cg::this_thread_block()); + const int tile_global = + blockIdx.x * (blockDim.x / TILE_SIZE) + (threadIdx.x / TILE_SIZE); + if (tile_global >= key_count) { + return; + } + const K key = keys[tile_global]; + const int pos = nv::merlin::device::find_readonly_no_sync( + view, key, tile.thread_rank()); + const unsigned int vote = tile.ballot(pos >= 0); + int found_pos = -1; + if (vote) { + const int src_lane = __ffs(vote) - 1; + found_pos = tile.shfl(pos, src_lane); + } + if (tile.thread_rank() == 0) { + positions[tile_global] = found_pos; + founds[tile_global] = (found_pos >= 0); + key_matches[tile_global] = + test_util::key_matches(view, key, found_pos); + } +} + +template +__global__ void ReadonlyLookupInBucketKernel( + nv::merlin::device::HashTableDeviceView view, const K* keys, + int key_count, bool* founds, int* positions, bool* key_matches) { + cg::thread_block_tile tile = + cg::tiled_partition(cg::this_thread_block()); + const int tile_global = + blockIdx.x * (blockDim.x / TILE_SIZE) + (threadIdx.x / TILE_SIZE); + if (tile_global >= key_count) { + return; + } + const K key = keys[tile_global]; + uint32_t bucket_idx = 0; + uint32_t aligned_start = 0; + if (!nv::merlin::device::compute_bucket_index_and_aligned_start( + view, key, &bucket_idx, &aligned_start)) { + if (tile.thread_rank() == 0) { + positions[tile_global] = -1; + founds[tile_global] = false; + key_matches[tile_global] = false; + } + return; + } + auto* bucket = view.buckets + bucket_idx; + const int pos = + nv::merlin::device::find_readonly_no_sync_in_bucket( + bucket, key, aligned_start, + static_cast(view.bucket_max_size), tile.thread_rank()); + const unsigned int vote = tile.ballot(pos >= 0); + int found_pos = -1; + if (vote) { + const int src_lane = __ffs(vote) - 1; + found_pos = tile.shfl(pos, src_lane); + } + if (tile.thread_rank() == 0) { + positions[tile_global] = found_pos; + founds[tile_global] = (found_pos >= 0); + key_matches[tile_global] = + test_util::key_matches(view, key, found_pos); + } +} + +namespace { + +void RunReadonlyLookupTest(bool use_in_bucket) { + TableOptions options; + options.init_capacity = 1024; + options.max_capacity = 1024; + options.max_bucket_size = 128; + options.dim = kDim; + options.max_hbm_for_vectors = nv::merlin::GB(1); + options.reserved_key_start_bit = 2; + options.num_of_buckets_per_alloc = 1; + + Table table; + table.init(options); + + constexpr int insert_count = 4; + std::array h_keys = {1, 2, 3, 4}; + std::array h_values{}; + std::array h_scores{}; + for (int i = 0; i < insert_count; ++i) { + h_scores[i] = static_cast(i + 1); + for (size_t j = 0; j < kDim; ++j) { + h_values[i * kDim + j] = static_cast(h_keys[i] * 0.1f + j); + } + } + + K* d_keys = nullptr; + V* d_values = nullptr; + S* d_scores = nullptr; + CUDA_CHECK(cudaMalloc(&d_keys, sizeof(K) * insert_count)); + CUDA_CHECK(cudaMalloc(&d_values, sizeof(V) * insert_count * kDim)); + CUDA_CHECK(cudaMalloc(&d_scores, sizeof(S) * insert_count)); + CUDA_CHECK(cudaMemcpy(d_keys, h_keys.data(), sizeof(K) * insert_count, + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(d_values, h_values.data(), + sizeof(V) * insert_count * kDim, + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(d_scores, h_scores.data(), sizeof(S) * insert_count, + cudaMemcpyHostToDevice)); + + table.insert_or_assign(insert_count, d_keys, d_values, d_scores); + CUDA_CHECK(cudaDeviceSynchronize()); + + constexpr int query_count = 6; + std::array h_query = {1, 4, 5, 2, 42, 3}; + K* d_query = nullptr; + CUDA_CHECK(cudaMalloc(&d_query, sizeof(K) * query_count)); + CUDA_CHECK(cudaMemcpy(d_query, h_query.data(), sizeof(K) * query_count, + cudaMemcpyHostToDevice)); + + bool* d_founds = nullptr; + int* d_positions = nullptr; + bool* d_key_matches = nullptr; + CUDA_CHECK(cudaMalloc(&d_founds, sizeof(bool) * query_count)); + CUDA_CHECK(cudaMalloc(&d_positions, sizeof(int) * query_count)); + CUDA_CHECK(cudaMalloc(&d_key_matches, sizeof(bool) * query_count)); + CUDA_CHECK(cudaMemset(d_founds, 0, sizeof(bool) * query_count)); + CUDA_CHECK(cudaMemset(d_positions, 0xff, sizeof(int) * query_count)); + CUDA_CHECK(cudaMemset(d_key_matches, 0, sizeof(bool) * query_count)); + + auto view = nv::merlin::device::make_device_view(table); + constexpr int threads = 128; + constexpr int tiles_per_block = threads / kTileSize; + const int blocks = (query_count + tiles_per_block - 1) / tiles_per_block; + if (use_in_bucket) { + ReadonlyLookupInBucketKernel<<>>( + view, d_query, query_count, d_founds, d_positions, d_key_matches); + } else { + ReadonlyLookupKernel<<>>( + view, d_query, query_count, d_founds, d_positions, d_key_matches); + } + CUDA_CHECK(cudaDeviceSynchronize()); + + std::array h_founds{}; + std::array h_positions{}; + std::array h_key_matches{}; + CUDA_CHECK(cudaMemcpy(h_founds.data(), d_founds, sizeof(bool) * query_count, + cudaMemcpyDeviceToHost)); + CUDA_CHECK(cudaMemcpy(h_positions.data(), d_positions, + sizeof(int) * query_count, cudaMemcpyDeviceToHost)); + CUDA_CHECK(cudaMemcpy(h_key_matches.data(), d_key_matches, + sizeof(bool) * query_count, cudaMemcpyDeviceToHost)); + + EXPECT_TRUE(h_founds[0]); + EXPECT_TRUE(h_founds[1]); + EXPECT_FALSE(h_founds[2]); + EXPECT_TRUE(h_founds[3]); + EXPECT_FALSE(h_founds[4]); + EXPECT_TRUE(h_founds[5]); + + for (int i = 0; i < query_count; ++i) { + if (h_founds[i]) { + EXPECT_GE(h_positions[i], 0); + EXPECT_TRUE(h_key_matches[i]); + } else { + EXPECT_EQ(h_positions[i], -1); + EXPECT_FALSE(h_key_matches[i]); + } + } + + CUDA_CHECK(cudaFree(d_keys)); + CUDA_CHECK(cudaFree(d_values)); + CUDA_CHECK(cudaFree(d_scores)); + CUDA_CHECK(cudaFree(d_query)); + CUDA_CHECK(cudaFree(d_founds)); + CUDA_CHECK(cudaFree(d_positions)); + CUDA_CHECK(cudaFree(d_key_matches)); +} + +} // namespace + +TEST(HashTableDeviceApiTest, ReadonlyLookup) { RunReadonlyLookupTest(false); } + +TEST(HashTableDeviceApiTest, ReadonlyLookupInBucket) { + RunReadonlyLookupTest(true); +} diff --git a/tests/test_util.cuh b/tests/test_util.cuh index 73e59157..15ecb485 100644 --- a/tests/test_util.cuh +++ b/tests/test_util.cuh @@ -30,6 +30,7 @@ #include #include "merlin/utils.cuh" #include "merlin_hashtable.cuh" +#include "merlin_hashtable_device.cuh" #define UNEQUAL_EXPR(expr1, expr2) \ { \ @@ -45,6 +46,27 @@ namespace test_util { +template +__device__ __forceinline__ bool key_matches( + const nv::merlin::device::HashTableDeviceView& view, + const typename nv::merlin::device::HashTableDeviceView::key_type + key, + const int pos) { + if (pos < 0) { + return false; + } + uint32_t bucket_idx = 0; + uint32_t aligned_start = 0; + if (!nv::merlin::device::compute_bucket_index_and_aligned_start( + view, key, &bucket_idx, &aligned_start)) { + return false; + } + const K current_key = + *reinterpret_cast((view.buckets + bucket_idx)->keys(pos)); + return current_key == key; +} + template __global__ void host_nano_kernel(S* d_clk) { S mclk;