CuckooFilterMultiGPU.cuh

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <memory>
#include <numeric>
#include <thread>
#include <type_traits>
#include <vector>
#include "CuckooFilter.cuh"
#include "helpers.cuh"
 
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/host_vector.h>
#include <thrust/scatter.h>
 
#include <gossip.cuh>
#include <plan_parser.hpp>
 
namespace cuckoogpu {
 
template <typename Config>
class FilterMultiGPU {
   public:
    using T = typename Config::KeyType;
 
    struct Partitioner {
        size_t numGPUs;
 
        __host__ __device__ gossip::gpu_id_t operator()(const T& key) const {
            uint64_t hash = Filter<Config>::hash64(key);
            return static_cast<gossip::gpu_id_t>(hash % numGPUs);
        }
    };
 
    static constexpr float defaultMemoryFactor = 0.8f;
 
   private:
    size_t numGPUs;
    size_t capacityPerGPU;
    float memoryFactor;
    std::vector<Filter<Config>*> filters;
 
    gossip::context_t gossipContext;
    gossip::multisplit_t multisplit;
    gossip::all2all_t all2all;
    gossip::all2all_t all2allResults;
 
    // Pre-allocated per-GPU buffers for gossip operations
    std::vector<T*> srcBuffers;
    std::vector<T*> dstBuffers;
    std::vector<size_t> bufferCapacities;
 
    std::vector<bool*> resultSrcBuffers;
    std::vector<bool*> resultDstBuffers;
 
    size_t totalBufferCapacity;
 
    [[nodiscard]] std::vector<size_t> getGpuMemoryInfo() const {
        std::vector<size_t> freeMem(numGPUs);
        parallelForGPUs([&](size_t gpuId) {
            size_t free, total;
            CUDA_CALL(cudaMemGetInfo(&free, &total));
            freeMem[gpuId] = free;
        });
        return freeMem;
    }
 
    void allocateBuffers() {
        // Bytes per element: 2 key buffers (src + dst) + 2 result buffers (src + dst)
        const size_t bytesPerKey = 2 * sizeof(T) + 2 * sizeof(bool);
 
        totalBufferCapacity = 0;
 
        parallelForGPUs([&](size_t gpuId) {
            size_t freeMem, totalMem;
            CUDA_CALL(cudaMemGetInfo(&freeMem, &totalMem));
 
            // Calculate max keys this GPU can buffer
            auto availableBytes = static_cast<size_t>(freeMem * memoryFactor);
            size_t maxKeys = availableBytes / bytesPerKey;
 
            // Allocate key buffers
            CUDA_CALL(cudaMalloc(&srcBuffers[gpuId], maxKeys * sizeof(T)));
            CUDA_CALL(cudaMalloc(&dstBuffers[gpuId], maxKeys * sizeof(T)));
 
            // Allocate result buffers
            CUDA_CALL(cudaMalloc(&resultSrcBuffers[gpuId], maxKeys * sizeof(bool)));
            CUDA_CALL(cudaMalloc(&resultDstBuffers[gpuId], maxKeys * sizeof(bool)));
 
            bufferCapacities[gpuId] = maxKeys;
            totalBufferCapacity += maxKeys;
        });
    }
 
    void freeBuffers() {
        parallelForGPUs([&](size_t gpuId) {
            if (srcBuffers[gpuId]) {
                cudaFree(srcBuffers[gpuId]);
                srcBuffers[gpuId] = nullptr;
            }
            if (dstBuffers[gpuId]) {
                cudaFree(dstBuffers[gpuId]);
                dstBuffers[gpuId] = nullptr;
            }
            if (resultSrcBuffers[gpuId]) {
                cudaFree(resultSrcBuffers[gpuId]);
                resultSrcBuffers[gpuId] = nullptr;
            }
            if (resultDstBuffers[gpuId]) {
                cudaFree(resultDstBuffers[gpuId]);
                resultDstBuffers[gpuId] = nullptr;
            }
            bufferCapacities[gpuId] = 0;
        });
    }
 
    template <bool returnOccupied, bool hasOutput, typename FilterFunc>
    size_t executeOperation(const T* h_keys, size_t n, bool* h_output, FilterFunc filterOp) {
        if (n == 0) {
            return returnOccupied ? totalOccupiedSlots() : 0;
        }
 
        size_t processed = 0;
 
        while (processed < n) {
            size_t chunkSize = std::min(n - processed, totalBufferCapacity);
 
            // Distribute chunk proportionally based on each GPU's buffer capacity
            std::vector<size_t> inputLens(numGPUs);
            std::vector<size_t> inputOffsets(numGPUs + 1, 0);
 
            size_t remaining = chunkSize;
            for (size_t gpu = 0; gpu < numGPUs; ++gpu) {
                if (gpu == numGPUs - 1) {
                    inputLens[gpu] = std::min(remaining, bufferCapacities[gpu]);
                } else {
                    double proportion =
                        static_cast<double>(bufferCapacities[gpu]) / totalBufferCapacity;
                    inputLens[gpu] = std::min(
                        static_cast<size_t>(chunkSize * proportion), bufferCapacities[gpu]
                    );
                    inputLens[gpu] = std::min(inputLens[gpu], remaining);
                }
                remaining -= inputLens[gpu];
                inputOffsets[gpu + 1] = inputOffsets[gpu] + inputLens[gpu];
            }
 
            // Copy input data to source buffers on each GPU
            parallelForGPUs([&](size_t gpuId) {
                if (inputLens[gpuId] > 0) {
                    CUDA_CALL(cudaMemcpy(
                        srcBuffers[gpuId],
                        h_keys + processed + inputOffsets[gpuId],
                        inputLens[gpuId] * sizeof(T),
                        cudaMemcpyHostToDevice
                    ));
                }
            });
            gossipContext.sync_hard();
 
            // Partition keys by target GPU
            std::vector<std::vector<size_t>> partitionTable(numGPUs, std::vector<size_t>(numGPUs));
 
            Partitioner partitioner{numGPUs};
            multisplit.execAsync(
                srcBuffers,        // source pointers (per GPU)
                inputLens,         // source lengths (per GPU)
                dstBuffers,        // destination pointers (per GPU)
                bufferCapacities,  // destination capacities (per GPU)
                partitionTable,    // output: partition counts [src][dst]
                partitioner
            );
            multisplit.sync();
 
            std::swap(srcBuffers, dstBuffers);
 
            // Calculate how many keys each GPU will receive after all2all
            std::vector<size_t> recvCounts(numGPUs, 0);
            for (size_t dst = 0; dst < numGPUs; ++dst) {
                for (size_t src = 0; src < numGPUs; ++src) {
                    recvCounts[dst] += partitionTable[src][dst];
                }
            }
 
            // Shuffle partitioned keys to correct GPUs
            all2all.execAsync(
                srcBuffers,        // partitioned source data
                bufferCapacities,  // source buffer capacities
                dstBuffers,        // destination for received data
                bufferCapacities,  // destination buffer capacities
                partitionTable     // partition counts from multisplit
            );
            all2all.sync();
 
            // If no output is required, execute filter ops and continue
            if constexpr (!hasOutput) {
                parallelForGPUs([&](size_t gpuId) {
                    size_t localCount = recvCounts[gpuId];
                    if (localCount == 0) {
                        return;
                    }
                    auto stream = gossipContext.get_streams(gpuId)[0];
                    filterOp(filters[gpuId], dstBuffers[gpuId], nullptr, localCount, stream);
                });
                gossipContext.sync_all_streams();
            } else {
                // Transpose partitionTable in-place for reverse all-to-all
                for (size_t i = 0; i < numGPUs; ++i) {
                    for (size_t j = i + 1; j < numGPUs; ++j) {
                        std::swap(partitionTable[i][j], partitionTable[j][i]);
                    }
                }
 
                // Execute filter operations
                parallelForGPUs([&](size_t gpuId) {
                    size_t localCount = recvCounts[gpuId];
                    if (localCount == 0) {
                        return;
                    }
                    auto stream = gossipContext.get_streams(gpuId)[0];
                    filterOp(
                        filters[gpuId],
                        dstBuffers[gpuId],
                        resultSrcBuffers[gpuId],
                        localCount,
                        stream
                    );
                });
                gossipContext.sync_all_streams();
 
                all2allResults.execAsync(
                    resultSrcBuffers, recvCounts, resultDstBuffers, bufferCapacities, partitionTable
                );
                all2allResults.sync();
 
                parallelForGPUs([&](size_t gpuId) {
                    size_t localCount = inputLens[gpuId];
                    if (localCount == 0) {
                        return;
                    }
 
                    CUDA_CALL(cudaMemcpy(
                        h_output + processed + inputOffsets[gpuId],
                        resultDstBuffers[gpuId],
                        localCount * sizeof(bool),
                        cudaMemcpyDeviceToHost
                    ));
                });
                gossipContext.sync_hard();
            }
 
            processed += chunkSize;
        }
 
        return returnOccupied ? totalOccupiedSlots() : 0;
    }
 
   public:
    FilterMultiGPU(size_t numGPUs, size_t capacity, float memFactor = defaultMemoryFactor)
        : numGPUs(numGPUs),
          capacityPerGPU(static_cast<size_t>(SDIV(capacity, numGPUs) * 1.02)),
          memoryFactor(memFactor),
          gossipContext(numGPUs),
          multisplit(gossipContext),
          all2all(gossipContext, gossip::all2all::default_plan(numGPUs)),
          all2allResults(gossipContext, gossip::all2all::default_plan(numGPUs)),
          srcBuffers(numGPUs, nullptr),
          dstBuffers(numGPUs, nullptr),
          bufferCapacities(numGPUs, 0),
          resultSrcBuffers(numGPUs, nullptr),
          resultDstBuffers(numGPUs, nullptr),
          totalBufferCapacity(0) {
        assert(numGPUs > 0 && "Number of GPUs must be at least 1");
 
        filters.resize(numGPUs);
 
        for (size_t i = 0; i < numGPUs; ++i) {
            CUDA_CALL(cudaSetDevice(gossipContext.get_device_id(i)));
            Filter<Config>* filter;
            CUDA_CALL(cudaMallocManaged(&filter, sizeof(Filter<Config>)));
            new (filter) Filter<Config>(capacityPerGPU);
            filters[i] = filter;
        }
        gossipContext.sync_hard();
 
        allocateBuffers();
    }
 
    FilterMultiGPU(
        size_t numGPUs,
        size_t capacity,
        const char* transferPlanPath,
        float memFactor = defaultMemoryFactor
    )
        : numGPUs(numGPUs),
          capacityPerGPU(static_cast<size_t>(SDIV(capacity, numGPUs) * 1.02)),
          memoryFactor(memFactor),
          gossipContext(numGPUs),
          multisplit(gossipContext),
          all2all(
              gossipContext,
              [&]() {
                  auto plan = parse_plan(transferPlanPath);
                  if (plan.num_gpus() == 0) {
                      return gossip::all2all::default_plan(numGPUs);
                  }
                  return plan;
              }()
          ),
          all2allResults(
              gossipContext,
              [&]() {
                  auto plan = parse_plan(transferPlanPath);
                  if (plan.num_gpus() == 0) {
                      return gossip::all2all::default_plan(numGPUs);
                  }
                  return plan;
              }()
          ),
          srcBuffers(numGPUs, nullptr),
          dstBuffers(numGPUs, nullptr),
          bufferCapacities(numGPUs, 0),
          resultSrcBuffers(numGPUs, nullptr),
          resultDstBuffers(numGPUs, nullptr),
          totalBufferCapacity(0) {
        assert(numGPUs > 0 && "Number of GPUs must be at least 1");
 
        filters.resize(numGPUs);
 
        for (size_t i = 0; i < numGPUs; ++i) {
            CUDA_CALL(cudaSetDevice(gossipContext.get_device_id(i)));
            Filter<Config>* filter;
            CUDA_CALL(cudaMallocManaged(&filter, sizeof(Filter<Config>)));
            new (filter) Filter<Config>(capacityPerGPU);
            filters[i] = filter;
        }
        gossipContext.sync_hard();
 
        allocateBuffers();
    }
 
    ~FilterMultiGPU() {
        freeBuffers();
        for (size_t i = 0; i < numGPUs; ++i) {
            CUDA_CALL(cudaSetDevice(gossipContext.get_device_id(i)));
            filters[i]->~Filter<Config>();
            CUDA_CALL(cudaFree(filters[i]));
        }
    }
 
    FilterMultiGPU(const FilterMultiGPU&) = delete;
    FilterMultiGPU& operator=(const FilterMultiGPU&) = delete;
 
    size_t insertMany(const T* h_keys, size_t n, bool* h_output = nullptr) {
        if (h_output) {
            return executeOperation<true, true>(
                h_keys,
                n,
                h_output,
                [](Filter<Config>* filter,
                   const T* keys,
                   bool* results,
                   size_t count,
                   cudaStream_t stream) { filter->insertMany(keys, count, results, stream); }
            );
        } else {
            return executeOperation<true, false>(
                h_keys,
                n,
                nullptr,
                [](Filter<Config>* filter,
                   const T* keys,
                   bool* /*unused results*/,
                   size_t count,
                   cudaStream_t stream) { filter->insertMany(keys, count, nullptr, stream); }
            );
        }
    }
 
    void containsMany(const T* h_keys, size_t n, bool* h_output) {
        executeOperation<false, true>(
            h_keys,
            n,
            h_output,
            [](Filter<Config>* filter,
               const T* keys,
               bool* results,
               size_t count,
               cudaStream_t stream) { filter->containsMany(keys, count, results, stream); }
        );
    }
 
    size_t deleteMany(const T* h_keys, size_t n, bool* h_output = nullptr) {
        if (h_output) {
            return executeOperation<true, true>(
                h_keys,
                n,
                h_output,
                [](Filter<Config>* filter,
                   const T* keys,
                   bool* results,
                   size_t count,
                   cudaStream_t stream) { filter->deleteMany(keys, count, results, stream); }
            );
        } else {
            return executeOperation<true, false>(
                h_keys,
                n,
                nullptr,
                [](Filter<Config>* filter,
                   const T* keys,
                   bool* /*unused results*/,
                   size_t count,
                   cudaStream_t stream) { filter->deleteMany(keys, count, nullptr, stream); }
            );
        }
    }
 
    [[nodiscard]] float loadFactor() const {
        return static_cast<float>(totalOccupiedSlots()) / static_cast<float>(totalCapacity());
    }
 
    template <typename Func>
    void parallelForGPUs(Func func) const {
        std::vector<std::thread> threads;
        for (size_t i = 0; i < numGPUs; ++i) {
            threads.emplace_back([=, this]() {
                CUDA_CALL(cudaSetDevice(gossipContext.get_device_id(i)));
                func(i);
            });
        }
 
        for (auto& t : threads) {
            t.join();
        }
    }
 
    void synchronizeAllGPUs() {
        gossipContext.sync_all_streams();
    }
 
    [[nodiscard]] size_t totalOccupiedSlots() const {
        std::atomic<size_t> total(0);
        parallelForGPUs([&](size_t i) {
            total.fetch_add(filters[i]->occupiedSlots(), std::memory_order_relaxed);
        });
 
        return total.load();
    }
 
    void clear() {
        parallelForGPUs([&](size_t i) { filters[i]->clear(); });
    }
 
    [[nodiscard]] size_t totalCapacity() const {
        std::atomic<size_t> total(0);
        parallelForGPUs([&](size_t i) {
            total.fetch_add(filters[i]->capacity(), std::memory_order_relaxed);
        });
        return total.load();
    }
 
    [[nodiscard]] size_t sizeInBytes() const {
        std::atomic<size_t> total(0);
        parallelForGPUs([&](size_t i) {
            total.fetch_add(filters[i]->sizeInBytes(), std::memory_order_relaxed);
        });
        return total.load();
    }
 
    size_t insertMany(const thrust::host_vector<T>& h_keys, thrust::host_vector<bool>& h_output) {
        h_output.resize(h_keys.size());
        return insertMany(
            thrust::raw_pointer_cast(h_keys.data()),
            h_keys.size(),
            thrust::raw_pointer_cast(h_output.data())
        );
    }
 
    size_t
    insertMany(const thrust::host_vector<T>& h_keys, thrust::host_vector<uint8_t>& h_output) {
        h_output.resize(h_keys.size());
        return insertMany(
            thrust::raw_pointer_cast(h_keys.data()),
            h_keys.size(),
            reinterpret_cast<bool*>(thrust::raw_pointer_cast(h_output.data()))
        );
    }
 
    size_t insertMany(const thrust::host_vector<T>& h_keys) {
        return insertMany(thrust::raw_pointer_cast(h_keys.data()), h_keys.size(), nullptr);
    }
 
    void containsMany(const thrust::host_vector<T>& h_keys, thrust::host_vector<bool>& h_output) {
        h_output.resize(h_keys.size());
        containsMany(
            thrust::raw_pointer_cast(h_keys.data()),
            h_keys.size(),
            thrust::raw_pointer_cast(h_output.data())
        );
    }
 
    void
    containsMany(const thrust::host_vector<T>& h_keys, thrust::host_vector<uint8_t>& h_output) {
        h_output.resize(h_keys.size());
        containsMany(
            thrust::raw_pointer_cast(h_keys.data()),
            h_keys.size(),
            reinterpret_cast<bool*>(thrust::raw_pointer_cast(h_output.data()))
        );
    }
 
    size_t deleteMany(const thrust::host_vector<T>& h_keys, thrust::host_vector<bool>& h_output) {
        h_output.resize(h_keys.size());
        return deleteMany(
            thrust::raw_pointer_cast(h_keys.data()),
            h_keys.size(),
            thrust::raw_pointer_cast(h_output.data())
        );
    }
 
    size_t
    deleteMany(const thrust::host_vector<T>& h_keys, thrust::host_vector<uint8_t>& h_output) {
        h_output.resize(h_keys.size());
        return deleteMany(
            thrust::raw_pointer_cast(h_keys.data()),
            h_keys.size(),
            reinterpret_cast<bool*>(thrust::raw_pointer_cast(h_output.data()))
        );
    }
 
    size_t deleteMany(const thrust::host_vector<T>& h_keys) {
        return deleteMany(thrust::raw_pointer_cast(h_keys.data()), h_keys.size(), nullptr);
    }
};
 
}  // namespace cuckoogpu