Use pre-allocated device memory for cuco hash tables

[xSke]

I've been trying to use this library in a context of PyCUDA/CuPy/etc, where I can compile device kernels just fine, but there's no way to interact with the host-side C++ functions (at least without a large amount of glue code...).

Is there a way to set up a static map using some block of pre-allocated memory (either a global on the device or passed in as a pointer) and then have all the operations from there (including initialization of some kind) take place in a device kernel I can call from Python?

I figure there'd be limitations here (like not being able to grow that memory), but it seems like a fairly common use case I can't quite work out how to use...

[PointKernel]

@xSke Thanks for bringing up this request. It appears I overlooked this GitHub issue until now. As far as I can tell, this is a completely valid feature request. I may be overlooking something, but it seems that what you're asking for is already feasible by passing a device pointer to static_map_ref. Specifically, please take a look at this example:
https://github.com/NVIDIA/cuCollections/blob/dev/examples/static_set/device_subsets_example.cu
At line 170, it shows how to create a non-own static_map_ref on the device using a device pointer and size.

Let me know if that addresses your use case or if there's something else you need to make it work.

[xSke]

Yes, it does! I figured it out a while after making this issue and forgot to make a reply.

It wasn't easy to figure out what types needed to be created where (especially as someone not the most proficient in heavily templated C++) so maybe another example would be nice to have - in particular one where the map is created entirely in initializers outside of a host function.

This is the example I came up with myself, if anyone else stumbles upon this issue, but it may be nice with a "cleaner" version:

#include <cuco/static_map.cuh>
#include <cuco/static_set.cuh>
#include <cuco/static_multimap.cuh>
#include <stdint.h>
#include <stdio.h>
#include <cuda/std/inplace_vector>

using Key   = int;
using Value = uint32_t;

Key constexpr empty_key_sentinel     = -1;
Value constexpr empty_value_sentinel = 0xffffffff;

std::size_t constexpr num_keys = 50'000;

auto constexpr load_factor = 0.5;
std::size_t const capacity = std::ceil(num_keys / load_factor);

auto constexpr bucket_size = 1;
auto constexpr cg_size = 1;
using probing_scheme_type = cuco::linear_probing<cg_size, cuco::default_hash_function<Key>>;

using extent_type = cuco::extent<size_t, capacity>;
using map_type = cuco::experimental::static_multimap<Key, Value, extent_type, cuda::thread_scope_block, cuda::std::equal_to<Key>,
                                    probing_scheme_type,
                                    cuco::cuda_allocator<Key>,
                                    cuco::storage<bucket_size>>;
using map_ref_type = map_type::ref_type<>;

using set_type = cuco::static_set<Value, extent_type>;

__device__ auto constexpr valid_extent =
    cuco::make_valid_extent<probing_scheme_type, cuco::storage<bucket_size>>(extent_type{});

__device__ cuda::std::array<typename map_ref_type::storage_ref_type::value_type, valid_extent.value()> arr;
__constant__ static size_t arr_size = sizeof(arr); // can read this from host code

__global__ void kern_insert() {
    auto storage = map_ref_type::storage_ref_type(valid_extent, arr.data());
    auto map = map_ref_type(cuco::empty_key{empty_key_sentinel}, cuco::empty_value{empty_value_sentinel}, {}, {}, {}, storage);

    auto id = threadIdx.x + blockIdx.x; 

    auto map_ref = map.rebind_operators(cuco::insert);
    map_ref.insert(cuco::pair(id, id * 2));

}

__global__ void kern_get() {
    auto storage = map_ref_type::storage_ref_type(valid_extent, arr.data());
    auto map = map_ref_type(cuco::empty_key{empty_key_sentinel}, cuco::empty_value{empty_value_sentinel}, {}, {}, {}, storage);

    auto id = threadIdx.x + blockIdx.x; 

    auto map2_ref = map.rebind_operators(cuco::find);
    auto res = map2_ref.find(id);

    if (res != map2_ref.end()) {
        printf("%d -> %d\n", res->first, res->second);
    }

}

__global__ void init() {
    auto storage = map_ref_type::storage_ref_type(valid_extent, arr.data());
    auto map = map_ref_type(cuco::empty_key{empty_key_sentinel}, cuco::empty_value{empty_value_sentinel}, {}, {}, {}, storage);
    auto const block = cooperative_groups::this_thread_block();

    map.initialize(block);
    block.sync();
}


int main() {
    init<<<1, 128>>>();
    cudaDeviceSynchronize();
    kern_insert<<<2, 32>>>();
    cudaDeviceSynchronize();
    kern_get<<<2, 32>>>();
    cudaDeviceSynchronize();
}

Here I've declared the data as a big static cuda::std::array, but if you were calling from eg. PyCUDA you could just as easily pass a raw pointer to arr into the kernels as a parameter (allocated using the size from arr_size, which I think you can read out from constant memory in PyCUDA?) and it should still work.

Thanks again!

[PointKernel]

It wasn't easy to figure out what types needed to be created where

Absolutely agree. We've been aiming to offer a concise interface similar to the STL, but that's challenging since STL doesn't deal with parallel execution or provide tuning options like we do. If you're okay with it, I'd like to add a new example based on your code so future users can benefit from it directly. Really appreciate you putting this example together!