ParallelZipTrie.cuh

Go to the documentation of this file.

 
#pragma once
 
#include "BitString.cuh" // Requires BitString definition with par_k_compare
 
#include <compare>   // For std::strong_ordering
 
#include "ZipTrie.hpp" // Base class definition
 
#include "cuda_utils.cuh" // CUDA utility functions like alloc_to_device, device_free
 
template <typename CHAR_T, bool MEMORY_EFFICIENT, typename RANK_T = GeometricRank, unsigned CHAR_SIZE_BITS = sizeof(CHAR_T) * 8>
class ParallelZipTrie : public ZipTrie<CHAR_T, MEMORY_EFFICIENT, RANK_T, CHAR_SIZE_BITS>
{
public:
    using ZT = ZipTrie<CHAR_T, MEMORY_EFFICIENT, RANK_T, CHAR_SIZE_BITS>;
 
    using LCP_T = typename ZT::LCP_T;
    using KEY_T = typename ZT::KEY_T;
 
    ParallelZipTrie(unsigned max_size, unsigned max_lcp_length);
 
    ~ParallelZipTrie();
 
    using ZT::size;
 
    void insert(const KEY_T* key) noexcept;
 
    // bool remove(const KEY_T& key) noexcept;
 
    using SearchResults = typename ZT::SearchResults;
 
    SearchResults search(const KEY_T* key) const noexcept override; // Mark as override
 
protected:
    // --- Inherited members from ZipTrie ---
    using ZT::_root_index;
    using ZT::_max_lcp_length;
    using ZT::_log_max_size;
    using ZT::_buckets;
 
    using ZT::NULLPTR;
 
    // --- Inherited types/structs from ZipTrie ---
    using AncestorLCPs = typename ZT::AncestorLCPs;
    using Bucket = typename ZT::Bucket;
    using ComparisonResult = typename ZT::ComparisonResult;
 
    // --- Inherited methods from ZipTrie ---
    using ZT::get_memory_efficient_lcp;
    using ZT::search_recursive;
    using ZT::insert_recursive;
    using ZT::k_compare_prefix_check;
    using ZT::convert_lcp;
 
private:
    uintmax_t* d_a;
    uintmax_t* d_largeblock;
 
    inline ComparisonResult k_compare(const KEY_T* x, const Bucket& v, const AncestorLCPs& ancestor_lcps, size_t& comparison_size, size_t& max_copied) const noexcept;
};
 
//-----------------------------------------------------------------------------
// ParallelZipTrie Method Implementations
//-----------------------------------------------------------------------------
 
template<typename CHAR_T, bool MEMORY_EFFICIENT, typename RANK_T, unsigned CHAR_SIZE_BITS>
ParallelZipTrie<CHAR_T, MEMORY_EFFICIENT, RANK_T, CHAR_SIZE_BITS>::ParallelZipTrie(unsigned max_size, unsigned max_lcp_length):
    ZT(max_size, max_lcp_length) // Call base class constructor
{
    // Calculate required device buffer size in words based on max LCP length.
    size_t max_lcp_length_words = (max_lcp_length + KEY_T::ALPHA - 1) / KEY_T::ALPHA;
 
    // Allocate the primary device buffer.
    d_a = alloc_to_device<uintmax_t>(max_lcp_length_words);
    // Allocate the auxiliary large block buffer (size calculation might depend on the parallel algorithm used).
    // Use large block device allocation for temporary storage in parallel operations
    d_largeblock = alloc_large_block_to_device(max_lcp_length_words);
}
 
template<typename CHAR_T, bool MEMORY_EFFICIENT, typename RANK_T, unsigned CHAR_SIZE_BITS>
ParallelZipTrie<CHAR_T, MEMORY_EFFICIENT, RANK_T, CHAR_SIZE_BITS>::~ParallelZipTrie()
{
    // Free the allocated device memory buffers.
    device_free(d_a);
    device_free(d_largeblock);
}
 
template<typename CHAR_T, bool MEMORY_EFFICIENT, typename RANK_T, unsigned CHAR_SIZE_BITS>
typename ParallelZipTrie<CHAR_T, MEMORY_EFFICIENT, RANK_T, CHAR_SIZE_BITS>::ComparisonResult ParallelZipTrie<CHAR_T, MEMORY_EFFICIENT, RANK_T, CHAR_SIZE_BITS>::k_compare(const KEY_T* x, const Bucket& v, const AncestorLCPs& ancestor_lcps, size_t& comparison_size, size_t& max_copied) const noexcept
{
    size_t start_lcp_val = 0; // Initialize LCP offset for comparison.
 
    if (auto prefix_result = k_compare_prefix_check(v, ancestor_lcps, start_lcp_val))
    {
        return *prefix_result; // Return early if comparison decided by LCPs.
    }
 
    // Initialize comparison result and final LCP with the value from prefix check.
    std::strong_ordering final_comparison = std::strong_ordering::equal;
    size_t final_lcp_val = start_lcp_val;
 
    // Iteratively call parallel compare until a mismatch is found or keys are fully compared.
    while (true)
    {
        // Call the parallel comparison function provided by the KEY_T (BitString) class.
        // Pass the keys, current LCP offset, comparison block size, device buffers, and copy tracker.
        auto [current_comparison, next_lcp] = x->par_k_compare(
            *v.key, final_lcp_val, comparison_size,
            d_a, d_largeblock, max_copied);
 
        // Update the overall comparison result and the total LCP found so far.
        final_comparison = current_comparison;
        final_lcp_val = next_lcp;
 
        // If the current block compared equal...
        if (final_comparison == std::strong_ordering::equal)
        {
            // Check if we have compared the entire length of both keys.
            // If keys are fully compared and equal, break the loop.
            if (final_lcp_val == x->size() && final_lcp_val == v.key->size())
            {
                break;
            }
 
            // Otherwise, keys matched so far but aren't finished.
            // Increase the comparison size (e.g., double it) for the next iteration to potentially speed up.
            comparison_size *= 2;
            // Continue to the next iteration to compare the next block.
            continue;
        }
 
        // Mismatch found (comparison is less or greater).
        // Reduce the comparison size for future comparisons (adaptive strategy).
        // Ensure it doesn't go below a minimum threshold.
        comparison_size = std::max(KEY_T::MIN_PAR_COMPARE_CHAR_SIZE, comparison_size / 2);
        // Break the loop as the final comparison result is determined.
        break;
    }
 
    // Convert the final exact LCP length to the appropriate LCP_T format and return the result.
    return { final_comparison, convert_lcp(final_lcp_val) };
}
 
template<typename CHAR_T, bool MEMORY_EFFICIENT, typename RANK_T, unsigned CHAR_SIZE_BITS>
typename ParallelZipTrie<CHAR_T, MEMORY_EFFICIENT, RANK_T, CHAR_SIZE_BITS>::SearchResults ParallelZipTrie<CHAR_T, MEMORY_EFFICIENT, RANK_T, CHAR_SIZE_BITS>::search(const KEY_T* key) const noexcept
{
    // Initialize state variables for the parallel comparison process.
    size_t comparison_size = KEY_T::MIN_PAR_COMPARE_CHAR_SIZE;
    size_t max_copied = 0; // Tracks device memory copy status/optimization.
 
    // Define a lambda function that captures the current state and calls the parallel k_compare method.
    // This lambda will be passed to the base class's search_recursive helper.
    auto parallel_compare = [&](const KEY_T* k, const Bucket& v, const AncestorLCPs& ancestor_lcps)
        {
            // Note: comparison_size and max_copied are captured by reference, allowing k_compare to modify them.
            return k_compare(k, v, ancestor_lcps, comparison_size, max_copied);
        };
 
    // Call the inherited recursive search function, passing the key and the parallel comparison lambda.
    return search_recursive(key, parallel_compare);
}
 
template <typename CHAR_T, bool MEMORY_EFFICIENT, typename RANK_T, unsigned CHAR_SIZE_BITS>
void ParallelZipTrie<CHAR_T, MEMORY_EFFICIENT, RANK_T, CHAR_SIZE_BITS>::insert(const KEY_T* key) noexcept
{
    // Initialize state variables for the parallel comparison process.
    size_t comparison_size = KEY_T::MIN_PAR_COMPARE_CHAR_SIZE;
    size_t max_copied = 0;
 
    // Add a new bucket for the key with a random rank to the storage vector.
    _buckets.push_back({ key, RANK_T::get_random() });
    unsigned new_node_index = _buckets.size() - 1;
 
    // Define the parallel comparison lambda, capturing state by reference.
    auto parallel_compare = [&](const KEY_T* k, const Bucket& v, const AncestorLCPs& ancestor_lcps)
        {
            return k_compare(k, v, ancestor_lcps, comparison_size, max_copied);
        };
 
    // Call the inherited recursive insert function to place the node and perform zip operations.
    // Pass the new node details, initial empty ancestor LCPs, and the parallel comparison lambda.
    // Update the root index, as it might change due to zip operations.
    _root_index = insert_recursive(&_buckets[new_node_index], new_node_index, _root_index, {}, parallel_compare);
}