CuckooFilter.cuh

Go to the documentation of this file.

#pragma once
 
#include <thrust/device_vector.h>
#include <cmath>
#include <cstdint>
#include <ctime>
#include <cub/cub.cuh>
#include <cuda/std/atomic>
#include <cuda/std/cstddef>
#include <cuda/std/cstdint>
#include <iostream>
#include <vector>
#include "bucket_policies.cuh"
#include "hashutil.cuh"
#include "helpers.cuh"
 
namespace cuckoogpu {
 
enum class EvictionPolicy {
    BFS,  
    DFS   
};
 
template <
    typename T,
    size_t bitsPerTag_,
    size_t maxEvictions_ = 500,
    size_t blockSize_ = 256,
    size_t bucketSize_ = 16,
    template <typename, typename, size_t, size_t> class AltBucketPolicy_ = XorAltBucketPolicy,
    EvictionPolicy evictionPolicy_ = EvictionPolicy::BFS,
    typename WordType_ = uint64_t>
struct Config {
    using KeyType = T;
    static constexpr size_t bitsPerTag = bitsPerTag_;
    static constexpr size_t maxEvictions = maxEvictions_;
    static constexpr size_t blockSize = blockSize_;
    static constexpr size_t bucketSize = bucketSize_;
    static constexpr EvictionPolicy evictionPolicy = evictionPolicy_;
 
    using TagType = typename std::conditional<
        bitsPerTag <= 8,
        uint8_t,
        typename std::conditional<bitsPerTag <= 16, uint16_t, uint32_t>::type>::type;
 
    using WordType = WordType_;
    static_assert(
        std::is_same_v<WordType, uint32_t> || std::is_same_v<WordType, uint64_t>,
        "WordType must be uint32_t or uint64_t"
    );
    static_assert(sizeof(TagType) <= sizeof(WordType), "TagType must fit within WordType");
 
    using AltBucketPolicy = AltBucketPolicy_<KeyType, TagType, bitsPerTag, bucketSize_>;
};
 
template <typename Config>
class Filter;
 
namespace detail {
 
template <typename Config>
__global__ void insertKernel(
    const typename Config::KeyType* keys,
    bool* output,
    size_t n,
    Filter<Config>* filter,
    uint32_t* evictionAttempts
);
 
template <typename Config>
__global__ void insertKernelSorted(
    const typename Filter<Config>::PackedTagType* packedTags,
    bool* output,
    size_t n,
    Filter<Config>* filter,
    uint32_t* evictionAttempts
);
 
template <typename Config>
__global__ void computePackedTagsKernel(
    const typename Config::KeyType* keys,
    typename Filter<Config>::PackedTagType* packedTags,
    size_t n,
    size_t numBuckets
);
 
template <typename Config>
__global__ void containsKernel(
    const typename Config::KeyType* keys,
    bool* output,
    size_t n,
    Filter<Config>* filter
);
 
template <typename Config>
__global__ void
deleteKernel(const typename Config::KeyType* keys, bool* output, size_t n, Filter<Config>* filter);
 
}  // namespace detail
 
template <typename Config>
struct Filter {
    using T = typename Config::KeyType;
    static constexpr size_t bitsPerTag = Config::bitsPerTag;
 
    using TagType = typename Config::TagType;
    using AltBucketPolicy = typename Config::AltBucketPolicy;
 
    static constexpr size_t tagEntryBytes = sizeof(TagType);
    static constexpr size_t bucketSize = Config::bucketSize;
 
    static constexpr size_t maxEvictions = Config::maxEvictions;
    static constexpr size_t blockSize = Config::blockSize;
    static_assert(
        bitsPerTag == 8 || bitsPerTag == 16 || bitsPerTag == 32,
        "The tag must be 8, 16 or 32 bits"
    );
 
    static_assert(detail::powerOfTwo(bucketSize), "Bucket size must be a power of 2");
 
    using PackedTagType = typename std::conditional<bitsPerTag <= 8, uint32_t, uint64_t>::type;
 
    struct PackedTag {
        PackedTagType value;
 
        // Lower bits = fingerprint
        // Upper bits = bucket index
        static constexpr size_t fpBits = bitsPerTag;
        static constexpr size_t totalBits = sizeof(PackedTagType) * 8;
        static constexpr size_t bucketIdxBits = totalBits - fpBits;
 
        static_assert(fpBits < totalBits, "fpBits must leave at least some bits for bucketIdx");
 
        static constexpr PackedTagType fpMask = PackedTagType((1ULL << fpBits) - 1ULL);
 
        static constexpr PackedTagType bucketIdxMask =
            PackedTagType(((1ULL << bucketIdxBits) - 1ULL) << fpBits);
 
        __host__ __device__ PackedTag() : value(0) {
        }
 
        __host__ __device__ PackedTag(TagType fp, uint64_t bucketIdx) : value(0) {
            setFingerprint(fp);
            setBucketIdx(bucketIdx);
        }
 
        __host__ __device__ TagType getFingerprint() const {
            return static_cast<TagType>(value & fpMask);
        }
 
        __host__ __device__ uint64_t getBucketIndex() const {
            return uint64_t((value & bucketIdxMask) >> fpBits);
        }
 
        __host__ __device__ void setFingerprint(TagType fp) {
            value = (value & ~fpMask) | (static_cast<PackedTagType>(fp) & fpMask);
        }
 
        __host__ __device__ void setBucketIdx(size_t bucketIdx) {
            PackedTagType v = static_cast<PackedTagType>(bucketIdx) << fpBits;
            value = (value & ~bucketIdxMask) | v;
        }
    };
 
    static constexpr TagType EMPTY = 0;
    static constexpr size_t fpMask = (1ULL << bitsPerTag) - 1;
 
    struct Bucket {
        static_assert(detail::powerOfTwo(bitsPerTag), "bitsPerTag must be a power of 2");
 
        using WordType = typename Config::WordType;
 
        static constexpr size_t tagsPerWord = sizeof(WordType) / sizeof(TagType);
        static_assert(tagsPerWord >= 1, "TagType must fit within WordType");
        static_assert(bucketSize % tagsPerWord == 0, "bucketSize must be divisible by tagsPerWord");
        static_assert(detail::powerOfTwo(tagsPerWord), "tagsPerWord must be a power of 2");
 
        static constexpr size_t wordCount = bucketSize / tagsPerWord;
        static_assert(detail::powerOfTwo(wordCount), "wordCount must be a power of 2");
 
        cuda::std::atomic<WordType> packedTags[wordCount];
 
        __host__ __device__ __forceinline__ TagType
        extractTag(WordType packed, size_t tagIdx) const {
            return static_cast<TagType>((packed >> (tagIdx * bitsPerTag)) & fpMask);
        }
 
        __host__ __device__ __forceinline__ WordType
        replaceTag(WordType packed, size_t tagIdx, TagType newTag) const {
            size_t shift = tagIdx * bitsPerTag;
            WordType cleared = packed & ~(static_cast<WordType>(fpMask) << shift);
            return cleared | (static_cast<WordType>(newTag) << shift);
        }
 
        template <size_t N>
        __device__ __forceinline__ static bool
        checkWords(const WordType (&loaded)[N], WordType replicatedTag) {
            _Pragma("unroll")
            for (size_t j = 0; j < N; ++j) {
                if (detail::hasZero<TagType, WordType>(loaded[j] ^ replicatedTag)) {
                    return true;
                }
            }
            return false;
        }
 
        template <size_t N>
        __device__ __forceinline__ void load128Bit(size_t startIdx, WordType (&out)[N]) const {
            static_assert(N == 2 || N == 4, "128-bit loads support 2 or 4 words");
            if constexpr (sizeof(WordType) == 4) {
                auto vec = __ldg(reinterpret_cast<const uint4*>(&packedTags[startIdx]));
                out[0] = vec.x;
                out[1] = vec.y;
                out[2] = vec.z;
                out[3] = vec.w;
            } else {
                auto vec = __ldg(reinterpret_cast<const ulonglong2*>(&packedTags[startIdx]));
                out[0] = vec.x;
                out[1] = vec.y;
            }
        }
 
        __device__ bool contains(TagType tag) const {
            const WordType replicatedTag = detail::replicateTag<TagType, WordType>(tag);
 
            // Scalar path
            if constexpr (wordCount == 1) {
                const auto packed = reinterpret_cast<const WordType&>(packedTags[0]);
                return detail::hasZero<TagType, WordType>(packed ^ replicatedTag);
            }
 
            const uint32_t startSlot = tag & (bucketSize - 1);
            const size_t startWordIdx = startSlot / tagsPerWord;
 
#if __CUDA_ARCH__ >= 1000 && !defined(CUCKOO_FILTER_DISABLE_256BIT_LOADS)
            // 256-bit load path
            constexpr size_t wordsPerLoad256 = (sizeof(WordType) == 4) ? 8 : 4;
            if constexpr (wordCount >= wordsPerLoad256) {
                constexpr size_t alignMask = wordsPerLoad256 - 1;
                const size_t startAlignedIdx = startWordIdx & ~alignMask;
 
                WordType loaded[wordsPerLoad256];
                for (size_t i = 0; i < wordCount / wordsPerLoad256; i++) {
                    const size_t idx = (startAlignedIdx + i * wordsPerLoad256) & (wordCount - 1);
                    detail::load256BitGlobalNC(
                        reinterpret_cast<const WordType*>(&packedTags[idx]), loaded
                    );
                    if (checkWords(loaded, replicatedTag)) {
                        return true;
                    }
                }
                return false;
            }
#endif
            // 128-bit load path
            constexpr size_t wordsPerLoad128 = (sizeof(WordType) == 4) ? 4 : 2;
            if constexpr (wordCount >= wordsPerLoad128) {
                constexpr size_t alignMask = wordsPerLoad128 - 1;
                const size_t startAlignedIdx = startWordIdx & ~alignMask;
 
                WordType loaded[wordsPerLoad128];
                for (size_t i = 0; i < wordCount / wordsPerLoad128; i++) {
                    const size_t idx = (startAlignedIdx + i * wordsPerLoad128) & (wordCount - 1);
                    load128Bit(idx, loaded);
                    if (checkWords(loaded, replicatedTag)) {
                        return true;
                    }
                }
                return false;
            }
 
            return false;
        }
    };
 
    size_t numBuckets;  
    Bucket* d_buckets;  
    cuda::std::atomic<size_t>*
        d_numOccupied{};  
 
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
    cuda::std::atomic<size_t>*
        d_numEvictions{};  
#endif
 
    size_t h_numOccupied = 0;  
 
    template <typename H>
    static __host__ __device__ uint64_t hash64(const H& key) {
        return AltBucketPolicy::hash64(key);
    }
 
    static __host__ __device__ cuda::std::tuple<size_t, size_t, TagType, TagType>
    getCandidateBucketsAndFPs(const T& key, size_t numBuckets) {
        return AltBucketPolicy::getCandidateBucketsAndFPs(key, numBuckets);
    }
 
    static __host__ __device__ size_t
    getAlternateBucket(size_t bucket, TagType fp, size_t numBuckets) {
        return AltBucketPolicy::getAlternateBucket(bucket, fp, numBuckets);
    }
 
    static __host__ __device__ cuda::std::tuple<size_t, TagType>
    getAlternateBucketWithNewFp(size_t bucket, TagType fp, size_t numBuckets) {
        if constexpr (AltBucketPolicy::usesChoiceBit) {
            return AltBucketPolicy::getAlternateBucketWithNewFp(bucket, fp, numBuckets);
        } else {
            return {AltBucketPolicy::getAlternateBucket(bucket, fp, numBuckets), fp};
        }
    }
 
    static size_t calculateNumBuckets(size_t capacity) {
        return AltBucketPolicy::calculateNumBuckets(capacity);
    }
 
    Filter(const Filter&) = delete;
    Filter& operator=(const Filter&) = delete;
 
    explicit Filter(size_t capacity) : numBuckets(calculateNumBuckets(capacity)) {
        CUDA_CALL(cudaMalloc(&d_buckets, numBuckets * sizeof(Bucket)));
        CUDA_CALL(cudaMalloc(&d_numOccupied, sizeof(cuda::std::atomic<size_t>)));
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
        CUDA_CALL(cudaMalloc(&d_numEvictions, sizeof(cuda::std::atomic<size_t>)));
#endif
 
        clear();
    }
 
    ~Filter() {
        if (d_buckets) {
            CUDA_CALL(cudaFree(d_buckets));
        }
        if (d_numOccupied) {
            CUDA_CALL(cudaFree(d_numOccupied));
        }
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
        if (d_numEvictions) {
            CUDA_CALL(cudaFree(d_numEvictions));
        }
#endif
    }
 
    size_t insertMany(
        const T* d_keys,
        const size_t n,
        bool* d_output = nullptr,
        cudaStream_t stream = {}
    ) {
        size_t numBlocks = SDIV(n, blockSize);
        detail::insertKernel<Config>
            <<<numBlocks, blockSize, 0, stream>>>(d_keys, d_output, n, this, nullptr);
 
        CUDA_CALL(cudaStreamSynchronize(stream));
 
        return occupiedSlots();
    }
 
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
    size_t insertManyWithEvictionCounts(
        const T* d_keys,
        const size_t n,
        uint32_t* d_evictionAttempts,
        bool* d_output = nullptr,
        cudaStream_t stream = {}
    ) {
        size_t numBlocks = SDIV(n, blockSize);
        detail::insertKernel<Config>
            <<<numBlocks, blockSize, 0, stream>>>(d_keys, d_output, n, this, d_evictionAttempts);
 
        CUDA_CALL(cudaStreamSynchronize(stream));
 
        return occupiedSlots();
    }
#endif
 
    size_t insertManySorted(
        const T* d_keys,
        const size_t n,
        bool* d_output = nullptr,
        cudaStream_t stream = {}
    ) {
        PackedTagType* d_packedTags;
 
        CUDA_CALL(cudaMallocAsync(&d_packedTags, n * sizeof(PackedTagType), stream));
 
        size_t numBlocks = SDIV(n, blockSize);
 
        detail::computePackedTagsKernel<Config>
            <<<numBlocks, blockSize, 0, stream>>>(d_keys, d_packedTags, n, numBuckets);
 
        void* d_tempStorage = nullptr;
        size_t tempStorageBytes = 0;
 
        cub::DeviceRadixSort::SortKeys(
            d_tempStorage,
            tempStorageBytes,
            d_packedTags,
            d_packedTags,
            n,
            0,
            sizeof(PackedTagType) * 8,
            stream
        );
 
        CUDA_CALL(cudaMallocAsync(&d_tempStorage, tempStorageBytes, stream));
 
        cub::DeviceRadixSort::SortKeys(
            d_tempStorage,
            tempStorageBytes,
            d_packedTags,
            d_packedTags,
            n,
            0,
            sizeof(PackedTagType) * 8,
            stream
        );
 
        CUDA_CALL(cudaFreeAsync(d_tempStorage, stream));
 
        detail::insertKernelSorted<Config>
            <<<numBlocks, blockSize, 0, stream>>>(d_packedTags, d_output, n, this, nullptr);
 
        CUDA_CALL(cudaFreeAsync(d_packedTags, stream));
        CUDA_CALL(cudaStreamSynchronize(stream));
 
        return occupiedSlots();
    }
 
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
    size_t insertManySortedWithEvictionCounts(
        const T* d_keys,
        const size_t n,
        uint32_t* d_evictionAttempts,
        bool* d_output = nullptr,
        cudaStream_t stream = {}
    ) {
        PackedTagType* d_packedTags;
 
        CUDA_CALL(cudaMallocAsync(&d_packedTags, n * sizeof(PackedTagType), stream));
 
        size_t numBlocks = SDIV(n, blockSize);
 
        detail::computePackedTagsKernel<Config>
            <<<numBlocks, blockSize, 0, stream>>>(d_keys, d_packedTags, n, numBuckets);
 
        void* d_tempStorage = nullptr;
        size_t tempStorageBytes = 0;
 
        cub::DeviceRadixSort::SortKeys(
            d_tempStorage,
            tempStorageBytes,
            d_packedTags,
            d_packedTags,
            n,
            0,
            sizeof(PackedTagType) * 8,
            stream
        );
 
        CUDA_CALL(cudaMallocAsync(&d_tempStorage, tempStorageBytes, stream));
 
        cub::DeviceRadixSort::SortKeys(
            d_tempStorage,
            tempStorageBytes,
            d_packedTags,
            d_packedTags,
            n,
            0,
            sizeof(PackedTagType) * 8,
            stream
        );
 
        CUDA_CALL(cudaFreeAsync(d_tempStorage, stream));
 
        detail::insertKernelSorted<Config><<<numBlocks, blockSize, 0, stream>>>(
            d_packedTags, d_output, n, this, d_evictionAttempts
        );
 
        CUDA_CALL(cudaFreeAsync(d_packedTags, stream));
        CUDA_CALL(cudaStreamSynchronize(stream));
 
        return occupiedSlots();
    }
#endif
 
    void containsMany(const T* d_keys, const size_t n, bool* d_output, cudaStream_t stream = {}) {
        size_t numBlocks = SDIV(n, blockSize);
        detail::containsKernel<Config>
            <<<numBlocks, blockSize, 0, stream>>>(d_keys, d_output, n, this);
 
        CUDA_CALL(cudaStreamSynchronize(stream));
    }
 
    size_t deleteMany(
        const T* d_keys,
        const size_t n,
        bool* d_output = nullptr,
        cudaStream_t stream = {}
    ) {
        size_t numBlocks = SDIV(n, blockSize);
        detail::deleteKernel<Config>
            <<<numBlocks, blockSize, 0, stream>>>(d_keys, d_output, n, this);
 
        CUDA_CALL(cudaStreamSynchronize(stream));
 
        return occupiedSlots();
    }
 
    size_t insertMany(
        const thrust::device_vector<T>& d_keys,
        thrust::device_vector<bool>& d_output,
        cudaStream_t stream = {}
    ) {
        if (d_output.size() != d_keys.size()) {
            d_output.resize(d_keys.size());
        }
        return insertMany(
            thrust::raw_pointer_cast(d_keys.data()),
            d_keys.size(),
            thrust::raw_pointer_cast(d_output.data()),
            stream
        );
    }
 
    size_t insertMany(
        const thrust::device_vector<T>& d_keys,
        thrust::device_vector<uint8_t>& d_output,
        cudaStream_t stream = {}
    ) {
        if (d_output.size() != d_keys.size()) {
            d_output.resize(d_keys.size());
        }
        return insertMany(
            thrust::raw_pointer_cast(d_keys.data()),
            d_keys.size(),
            reinterpret_cast<bool*>(thrust::raw_pointer_cast(d_output.data())),
            stream
        );
    }
 
    size_t insertMany(const thrust::device_vector<T>& d_keys, cudaStream_t stream = {}) {
        return insertMany(thrust::raw_pointer_cast(d_keys.data()), d_keys.size(), nullptr, stream);
    }
 
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
    size_t insertManyWithEvictionCounts(
        const thrust::device_vector<T>& d_keys,
        thrust::device_vector<uint32_t>& d_evictionAttempts,
        thrust::device_vector<uint8_t>* d_output = nullptr,
        cudaStream_t stream = {}
    ) {
        if (d_evictionAttempts.size() != d_keys.size()) {
            d_evictionAttempts.resize(d_keys.size());
        }
 
        bool* d_outputPtr = nullptr;
        if (d_output != nullptr) {
            if (d_output->size() != d_keys.size()) {
                d_output->resize(d_keys.size());
            }
            d_outputPtr = reinterpret_cast<bool*>(thrust::raw_pointer_cast(d_output->data()));
        }
 
        return insertManyWithEvictionCounts(
            thrust::raw_pointer_cast(d_keys.data()),
            d_keys.size(),
            thrust::raw_pointer_cast(d_evictionAttempts.data()),
            d_outputPtr,
            stream
        );
    }
#endif
 
    size_t insertManySorted(
        const thrust::device_vector<T>& d_keys,
        thrust::device_vector<bool>& d_output,
        cudaStream_t stream = {}
    ) {
        if (d_output.size() != d_keys.size()) {
            d_output.resize(d_keys.size());
        }
        return insertManySorted(
            thrust::raw_pointer_cast(d_keys.data()),
            d_keys.size(),
            thrust::raw_pointer_cast(d_output.data()),
            stream
        );
    }
 
    size_t insertManySorted(
        const thrust::device_vector<T>& d_keys,
        thrust::device_vector<uint8_t>& d_output,
        cudaStream_t stream = {}
    ) {
        if (d_output.size() != d_keys.size()) {
            d_output.resize(d_keys.size());
        }
        return insertManySorted(
            thrust::raw_pointer_cast(d_keys.data()),
            d_keys.size(),
            reinterpret_cast<bool*>(thrust::raw_pointer_cast(d_output.data())),
            stream
        );
    }
 
    size_t insertManySorted(const thrust::device_vector<T>& d_keys, cudaStream_t stream = {}) {
        return insertManySorted(
            thrust::raw_pointer_cast(d_keys.data()), d_keys.size(), nullptr, stream
        );
    }
 
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
    size_t insertManySortedWithEvictionCounts(
        const thrust::device_vector<T>& d_keys,
        thrust::device_vector<uint32_t>& d_evictionAttempts,
        thrust::device_vector<uint8_t>* d_output = nullptr,
        cudaStream_t stream = {}
    ) {
        if (d_evictionAttempts.size() != d_keys.size()) {
            d_evictionAttempts.resize(d_keys.size());
        }
 
        bool* d_outputPtr = nullptr;
        if (d_output != nullptr) {
            if (d_output->size() != d_keys.size()) {
                d_output->resize(d_keys.size());
            }
            d_outputPtr = reinterpret_cast<bool*>(thrust::raw_pointer_cast(d_output->data()));
        }
 
        return insertManySortedWithEvictionCounts(
            thrust::raw_pointer_cast(d_keys.data()),
            d_keys.size(),
            thrust::raw_pointer_cast(d_evictionAttempts.data()),
            d_outputPtr,
            stream
        );
    }
#endif
 
    void containsMany(
        const thrust::device_vector<T>& d_keys,
        thrust::device_vector<bool>& d_output,
        cudaStream_t stream = {}
    ) {
        if (d_output.size() != d_keys.size()) {
            d_output.resize(d_keys.size());
        }
        containsMany(
            thrust::raw_pointer_cast(d_keys.data()),
            d_keys.size(),
            thrust::raw_pointer_cast(d_output.data()),
            stream
        );
    }
 
    void containsMany(
        const thrust::device_vector<T>& d_keys,
        thrust::device_vector<uint8_t>& d_output,
        cudaStream_t stream = {}
    ) {
        if (d_output.size() != d_keys.size()) {
            d_output.resize(d_keys.size());
        }
        containsMany(
            thrust::raw_pointer_cast(d_keys.data()),
            d_keys.size(),
            reinterpret_cast<bool*>(thrust::raw_pointer_cast(d_output.data())),
            stream
        );
    }
 
    size_t deleteMany(
        const thrust::device_vector<T>& d_keys,
        thrust::device_vector<bool>& d_output,
        cudaStream_t stream = {}
    ) {
        if (d_output.size() != d_keys.size()) {
            d_output.resize(d_keys.size());
        }
        return deleteMany(
            thrust::raw_pointer_cast(d_keys.data()),
            d_keys.size(),
            thrust::raw_pointer_cast(d_output.data()),
            stream
        );
    }
 
    size_t deleteMany(
        const thrust::device_vector<T>& d_keys,
        thrust::device_vector<uint8_t>& d_output,
        cudaStream_t stream = {}
    ) {
        if (d_output.size() != d_keys.size()) {
            d_output.resize(d_keys.size());
        }
        return deleteMany(
            thrust::raw_pointer_cast(d_keys.data()),
            d_keys.size(),
            reinterpret_cast<bool*>(thrust::raw_pointer_cast(d_output.data())),
            stream
        );
    }
 
    size_t deleteMany(const thrust::device_vector<T>& d_keys, cudaStream_t stream = {}) {
        return deleteMany(thrust::raw_pointer_cast(d_keys.data()), d_keys.size(), nullptr, stream);
    }
 
    void clear() {
        CUDA_CALL(cudaMemset(d_buckets, 0, numBuckets * sizeof(Bucket)));
        CUDA_CALL(cudaMemset(d_numOccupied, 0, sizeof(cuda::std::atomic<size_t>)));
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
        CUDA_CALL(cudaMemset(d_numEvictions, 0, sizeof(cuda::std::atomic<size_t>)));
#endif
        h_numOccupied = 0;
    }
 
    [[nodiscard]] float loadFactor() {
        return static_cast<float>(occupiedSlots()) / (numBuckets * bucketSize);
    }
 
    size_t occupiedSlots() {
        CUDA_CALL(
            cudaMemcpy(&h_numOccupied, d_numOccupied, sizeof(size_t), cudaMemcpyDeviceToHost)
        );
        return h_numOccupied;
    }
 
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
    size_t evictionCount() {
        size_t count;
        CUDA_CALL(cudaMemcpy(&count, d_numEvictions, sizeof(size_t), cudaMemcpyDeviceToHost));
        return count;
    }
 
    void resetEvictionCount() {
        CUDA_CALL(cudaMemset(d_numEvictions, 0, sizeof(cuda::std::atomic<size_t>)));
    }
#endif
 
    size_t capacity() {
        return numBuckets * bucketSize;
    }
 
    [[nodiscard]] size_t getNumBuckets() const {
        return numBuckets;
    }
 
    [[nodiscard]] size_t sizeInBytes() const {
        return numBuckets * sizeof(Bucket);
    }
 
    size_t countOccupiedSlots() {
        std::vector<Bucket> h_buckets(numBuckets);
 
        CUDA_CALL(cudaMemcpy(
            h_buckets.data(), d_buckets, numBuckets * sizeof(Bucket), cudaMemcpyDeviceToHost
        ));
 
        size_t occupiedCount = 0;
 
        for (size_t bucketIdx = 0; bucketIdx < numBuckets; ++bucketIdx) {
            const Bucket& bucket = h_buckets[bucketIdx];
 
            for (size_t atomicIdx = 0; atomicIdx < Bucket::wordCount; ++atomicIdx) {
                uint64_t packed = reinterpret_cast<const uint64_t&>(bucket.packedTags[atomicIdx]);
 
                for (size_t tagIdx = 0; tagIdx < Bucket::tagsPerWord; ++tagIdx) {
                    auto tag = bucket.extractTag(packed, tagIdx);
 
                    if (tag != EMPTY) {
                        occupiedCount++;
                    }
                }
            }
        }
 
        return occupiedCount;
    }
 
    __device__ bool tryRemoveAtBucket(size_t bucketIdx, TagType tag) {
        Bucket& bucket = d_buckets[bucketIdx];
 
        const uint32_t startSlot = tag & (bucketSize - 1);
        const size_t startWord = startSlot / Bucket::tagsPerWord;
 
        using WordType = typename Bucket::WordType;
 
        for (size_t i = 0; i < Bucket::wordCount; ++i) {
            const size_t currIdx = (startWord + i) & (Bucket::wordCount - 1);
 
            while (true) {
                auto expected = bucket.packedTags[currIdx].load(cuda::memory_order_relaxed);
 
                WordType matchMask = detail::getZeroMask<TagType, WordType>(
                    expected ^ detail::replicateTag<TagType, WordType>(tag)
                );
 
                if (matchMask == 0) {
                    // No matching tags in this word
                    break;
                }
 
                // Find position of first matching tag
                int bitPos;
                if constexpr (sizeof(WordType) == 4) {
                    bitPos = __ffs(static_cast<int>(matchMask)) - 1;
                } else {
                    bitPos = __ffsll(static_cast<long long>(matchMask)) - 1;
                }
                size_t tagIdx = bitPos / bitsPerTag;
 
                auto desired = bucket.replaceTag(expected, tagIdx, EMPTY);
 
                if (bucket.packedTags[currIdx].compare_exchange_weak(
                        expected, desired, cuda::memory_order_relaxed, cuda::memory_order_relaxed
                    )) {
                    return true;
                }
                // CAS failed, retry with updated expected value
            }
        }
 
        return false;
    }
 
    __device__ bool tryInsertAtBucket(size_t bucketIdx, TagType tag) {
        Bucket& bucket = d_buckets[bucketIdx];
        const uint32_t startIdx = tag & (bucketSize - 1);
        const size_t startWord = startIdx / Bucket::tagsPerWord;
 
        using WordType = typename Bucket::WordType;
 
        for (size_t i = 0; i < Bucket::wordCount; ++i) {
            const size_t currWord = (startWord + i) & (Bucket::wordCount - 1);
            auto expected = bucket.packedTags[currWord].load(cuda::memory_order_relaxed);
 
            while (true) {
                WordType zeroMask = detail::getZeroMask<TagType, WordType>(expected);
 
                if (zeroMask == 0) {
                    // No empty slots in this word, move to next
                    break;
                }
 
                // Find position of first empty slot (returns 1-indexed bit position)
                int bitPos;
                if constexpr (sizeof(WordType) == 4) {
                    bitPos = __ffs(static_cast<int>(zeroMask)) - 1;
                } else {
                    bitPos = __ffsll(static_cast<long long>(zeroMask)) - 1;
                }
                size_t j = bitPos / bitsPerTag;
 
                auto desired = bucket.replaceTag(expected, j, tag);
 
                if (bucket.packedTags[currWord].compare_exchange_strong(
                        expected, desired, cuda::memory_order_relaxed, cuda::memory_order_relaxed
                    )) {
                    return true;
                }
            }
        }
        return false;
    }
 
    __device__ bool
    insertWithEvictionDFS(TagType fp, size_t startBucket, uint32_t* evictionAttempts = nullptr) {
        TagType currentFp = fp;
        size_t currentBucket = startBucket;
 
        for (size_t evictions = 0; evictions < maxEvictions; ++evictions) {
            auto evictSlot = (currentFp + (evictions + 1) * 0x9E3779B1UL) & (bucketSize - 1);
 
            size_t evictWord = evictSlot / Bucket::tagsPerWord;
            size_t evictTagIdx = evictSlot & (Bucket::tagsPerWord - 1);
 
            Bucket& bucket = d_buckets[currentBucket];
            auto expected = bucket.packedTags[evictWord].load(cuda::memory_order_relaxed);
            typename Bucket::WordType desired;
            TagType evictedFp;
 
            do {
                evictedFp = bucket.extractTag(expected, evictTagIdx);
                desired = bucket.replaceTag(expected, evictTagIdx, currentFp);
            } while (!bucket.packedTags[evictWord].compare_exchange_strong(
                expected, desired, cuda::memory_order_relaxed, cuda::memory_order_relaxed
            ));
 
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
            d_numEvictions->fetch_add(1, cuda::memory_order_relaxed);
            if (evictionAttempts != nullptr) {
                (*evictionAttempts)++;
            }
#endif
 
            currentFp = evictedFp;
            auto [altBucket, newFp] =
                getAlternateBucketWithNewFp(currentBucket, evictedFp, numBuckets);
            currentBucket = altBucket;
            currentFp = newFp;
 
            if (tryInsertAtBucket(currentBucket, currentFp)) {
                return true;
            }
        }
        return false;
    }
 
    __device__ bool
    insertWithEvictionBFS(TagType fp, size_t startBucket, uint32_t* evictionAttempts = nullptr) {
        constexpr size_t numCandidates = std::max(1UL, bucketSize / 2);
 
        TagType currentFp = fp;
        size_t currentBucket = startBucket;
 
        size_t evictions = 0;
        while (evictions < maxEvictions) {
            Bucket& bucket = d_buckets[currentBucket];
            size_t restartWord = 0;
            size_t restartTagIdx = 0;
 
            for (size_t i = 0; i < numCandidates; ++i) {
                size_t evictSlot = (currentFp + i * 0x9E3779B1UL + (evictions + 1) * 0x85EBCA77) &
                                   (bucketSize - 1);
                size_t evictWord = evictSlot / Bucket::tagsPerWord;
                size_t evictTagIdx = evictSlot & (Bucket::tagsPerWord - 1);
                restartWord = evictWord;
                restartTagIdx = evictTagIdx;
 
                auto packed = bucket.packedTags[evictWord].load(cuda::memory_order_relaxed);
                TagType candidateFp = bucket.extractTag(packed, evictTagIdx);
 
                if (candidateFp == EMPTY) {
                    if (tryInsertAtBucket(currentBucket, currentFp)) {
                        return true;
                    }
                    continue;
                }
 
                auto [altBucket, altFp] =
                    getAlternateBucketWithNewFp(currentBucket, candidateFp, numBuckets);
                if (tryInsertAtBucket(altBucket, altFp)) {
                    // Successfully inserted the evicted tag at its alternate location
                    // Now atomically swap in our tag at the original location
                    auto expected = bucket.packedTags[evictWord].load(cuda::memory_order_relaxed);
 
                    // Verify the tag is still there and try to replace it
                    if (bucket.extractTag(expected, evictTagIdx) == candidateFp) {
                        auto desired = bucket.replaceTag(expected, evictTagIdx, currentFp);
 
                        if (bucket.packedTags[evictWord].compare_exchange_strong(
                                expected,
                                desired,
                                cuda::memory_order_relaxed,
                                cuda::memory_order_relaxed
                            )) {
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
                            d_numEvictions->fetch_add(1, cuda::memory_order_relaxed);
                            if (evictionAttempts != nullptr) {
                                (*evictionAttempts)++;
                            }
#endif
                            return true;
                        }
                    }
 
                    // Failed to swap, clean up the tag we inserted to avoid duplicates
                    tryRemoveAtBucket(altBucket, candidateFp);
                }
            }
 
            // Evict the last scanned candidate and continue from its alternate location.
            auto expected = bucket.packedTags[restartWord].load(cuda::memory_order_relaxed);
            typename Bucket::WordType desired;
            TagType evictedFp;
 
            do {
                evictedFp = bucket.extractTag(expected, restartTagIdx);
                desired = bucket.replaceTag(expected, restartTagIdx, currentFp);
            } while (!bucket.packedTags[restartWord].compare_exchange_strong(
                expected, desired, cuda::memory_order_relaxed, cuda::memory_order_relaxed
            ));
 
            if (evictedFp == EMPTY) {
                return true;
            }
 
#ifdef CUCKOO_FILTER_COUNT_EVICTIONS
            d_numEvictions->fetch_add(1, cuda::memory_order_relaxed);
            if (evictionAttempts != nullptr) {
                (*evictionAttempts)++;
            }
#endif
 
            evictions++;
            auto [altBucket, altFp] =
                getAlternateBucketWithNewFp(currentBucket, evictedFp, numBuckets);
            currentBucket = altBucket;
            currentFp = altFp;
        }
 
        return false;
    }
 
    __device__ bool insert(const T& key, uint32_t* evictionAttempts = nullptr) {
        auto [i1, i2, fp1, fp2] = getCandidateBucketsAndFPs(key, numBuckets);
 
        // For all policies: fp1 is for bucket i1, fp2 is for bucket i2
        // For non-choice-bit policies, fp1 == fp2
        if (tryInsertAtBucket(i1, fp1) || tryInsertAtBucket(i2, fp2)) {
            return true;
        }
 
        // For eviction, use correct fingerprint for the starting bucket
        auto startBucket = (fp1 & 1) == 0 ? i1 : i2;
        TagType evictFp;
 
        if constexpr (AltBucketPolicy::usesChoiceBit) {
            evictFp = (fp1 & 1) == 0 ? fp1 : fp2;
        } else {
            evictFp = fp1;
        }
 
        if constexpr (Config::evictionPolicy == EvictionPolicy::BFS) {
            return insertWithEvictionBFS(evictFp, startBucket, evictionAttempts);
        } else if constexpr (Config::evictionPolicy == EvictionPolicy::DFS) {
            return insertWithEvictionDFS(evictFp, startBucket, evictionAttempts);
        } else {
            static_assert(
                Config::evictionPolicy == EvictionPolicy::DFS ||
                    Config::evictionPolicy == EvictionPolicy::BFS,
                "Unhandled eviction policy"
            );
        }
    }
 
    __device__ bool contains(const T& key) const {
        auto [i1, i2, fp1, fp2] = getCandidateBucketsAndFPs(key, numBuckets);
 
        // fp1 is for bucket i1, fp2 is for bucket i2
        // For non-choice-bit policies, fp1 == fp2
        return d_buckets[i1].contains(fp1) || d_buckets[i2].contains(fp2);
    }
 
    __device__ bool remove(const T& key) {
        auto [i1, i2, fp1, fp2] = getCandidateBucketsAndFPs(key, numBuckets);
 
        // fp1 is for bucket i1, fp2 is for bucket i2
        // For non-choice-bit policies, fp1 == fp2
        return tryRemoveAtBucket(i1, fp1) || tryRemoveAtBucket(i2, fp2);
    }
};
 
namespace detail {
 
template <typename Config>
__global__ void insertKernel(
    const typename Config::KeyType* keys,
    bool* output,
    size_t n,
    Filter<Config>* filter,
    uint32_t* evictionAttempts
) {
    using BlockReduce = cub::BlockReduce<int32_t, Config::blockSize>;
    __shared__ typename BlockReduce::TempStorage tempStorage;
 
    auto idx = globalThreadId();
 
    int32_t success = 0;
 
    if (idx < n) {
        uint32_t threadEvictions = 0;
        success = filter->insert(keys[idx], &threadEvictions);
 
        if (output != nullptr) {
            output[idx] = success;
        }
 
        if (evictionAttempts != nullptr) {
            evictionAttempts[idx] = threadEvictions;
        }
    }
 
    int32_t blockSuccessSum = BlockReduce(tempStorage).Sum(success);
    __syncthreads();
 
    if (threadIdx.x == 0) {
        if (blockSuccessSum > 0) {
            filter->d_numOccupied->fetch_add(blockSuccessSum, cuda::memory_order_relaxed);
        }
    }
}
 
template <typename Config>
__global__ void containsKernel(
    const typename Config::KeyType* keys,
    bool* output,
    size_t n,
    Filter<Config>* filter
) {
    auto idx = globalThreadId();
 
    if (idx < n) {
        output[idx] = filter->contains(keys[idx]);
    }
}
 
template <typename Config>
__global__ void
deleteKernel(const typename Config::KeyType* keys, bool* output, size_t n, Filter<Config>* filter) {
    using BlockReduce = cub::BlockReduce<int32_t, Config::blockSize>;
    __shared__ typename BlockReduce::TempStorage tempStorage;
 
    auto idx = globalThreadId();
 
    int32_t success = 0;
    if (idx < n) {
        success = filter->remove(keys[idx]);
 
        if (output != nullptr) {
            output[idx] = success;
        }
    }
 
    int32_t blockSum = BlockReduce(tempStorage).Sum(success);
 
    if (threadIdx.x == 0 && blockSum > 0) {
        filter->d_numOccupied->fetch_sub(blockSum, cuda::memory_order_relaxed);
    }
}
 
template <typename Config>
__global__ void computePackedTagsKernel(
    const typename Config::KeyType* keys,
    typename Filter<Config>::PackedTagType* packedTags,
    size_t n,
    size_t numBuckets
) {
    size_t idx = globalThreadId();
 
    if (idx >= n) {
        return;
    }
 
    using FilterType = Filter<Config>;
    using PackedTagType = typename FilterType::PackedTagType;
    constexpr size_t bitsPerTag = Config::bitsPerTag;
 
    typename Config::KeyType key = keys[idx];
    auto [i1, i2, fp1, fp2] = FilterType::getCandidateBucketsAndFPs(key, numBuckets);
 
    packedTags[idx] =
        (static_cast<PackedTagType>(i1) << bitsPerTag) | static_cast<PackedTagType>(fp1);
}
 
template <typename Config>
__global__ void insertKernelSorted(
    const typename Filter<Config>::PackedTagType* packedTags,
    bool* output,
    size_t n,
    Filter<Config>* filter,
    uint32_t* evictionAttempts
) {
    using BlockReduce = cub::BlockReduce<int, Config::blockSize>;
    __shared__ typename BlockReduce::TempStorage tempStorage;
 
    size_t idx = globalThreadId();
 
    using FilterType = Filter<Config>;
    using TagType = typename FilterType::TagType;
    using PackedTagType = typename FilterType::PackedTagType;
 
    constexpr size_t bitsPerTag = Config::bitsPerTag;
    constexpr TagType fpMask = (1ULL << bitsPerTag) - 1;
 
    int32_t success = 0;
    uint32_t threadEvictions = 0;
    if (idx < n) {
        PackedTagType packedTag = packedTags[idx];
        size_t primaryBucket = packedTag >> bitsPerTag;
        auto fp = static_cast<TagType>(packedTag & fpMask);
 
        if (filter->tryInsertAtBucket(primaryBucket, fp)) {
            success = 1;
        } else {
            auto [i2, fp2] =
                FilterType::getAlternateBucketWithNewFp(primaryBucket, fp, filter->numBuckets);
 
            if (filter->tryInsertAtBucket(i2, fp2)) {
                success = 1;
            } else {
                TagType evictFp;
                auto startBucket = (fp & 1) == 0 ? primaryBucket : i2;
 
                if constexpr (Config::AltBucketPolicy::usesChoiceBit) {
                    evictFp = (fp & 1) == 0 ? fp : fp2;
                } else {
                    evictFp = fp;
                }
 
                if constexpr (Config::evictionPolicy == EvictionPolicy::BFS) {
                    success = filter->insertWithEvictionBFS(evictFp, startBucket, &threadEvictions);
                } else if constexpr (Config::evictionPolicy == EvictionPolicy::DFS) {
                    success = filter->insertWithEvictionDFS(evictFp, startBucket, &threadEvictions);
                } else {
                    static_assert(
                        Config::evictionPolicy == EvictionPolicy::DFS ||
                            Config::evictionPolicy == EvictionPolicy::BFS,
                        "Unhandled eviction policy"
                    );
                }
            }
        }
 
        if (output != nullptr) {
            output[idx] = success;
        }
 
        if (evictionAttempts != nullptr) {
            evictionAttempts[idx] = threadEvictions;
        }
    }
 
    int32_t blockSum = BlockReduce(tempStorage).Sum(success);
 
    if (threadIdx.x == 0 && blockSum > 0) {
        filter->d_numOccupied->fetch_add(blockSum, cuda::memory_order_relaxed);
    }
}
 
}  // namespace detail
 
}  // namespace cuckoogpu