Compiler Explorer

Source code

#include <cuco/static_map.cuh>

#include <thrust/device_vector.h>

#include <vector>

#include <cooperative_groups.h>

template <typename MapRef>
__global__ void insert_into_maps(MapRef* maps, std::size_t num_maps) {
  auto constexpr cg_size = MapRef::cg_size;
  auto const tid = blockIdx.x * blockDim.x + threadIdx.x;
  auto const gid = tid / cg_size;

auto const block = cooperative_groups::this_thread_block();
  auto const tile  = cooperative_groups::tiled_partition<cg_size>(block);
  auto const insert_pair = cuco::pair{gid, gid};
  maps[gid % num_maps].insert(tile, insert_pair);
}

int main(void)
{
  using Key = int;
  using T   = int;
  using map_type = cuco::static_map<Key, T>;

// 10 maps; each with capacity 1000
  std::size_t constexpr capacity = 1000;
  std::size_t num_maps = 10;
  Key constexpr empty_key_sentinel = -1;
  T constexpr empty_value_sentinel = -1;

// stores the actual maps, i.e., owns the memory
  std::vector<map_type> maps;
  // device vector that holds non-owning references to each map
  thrust::device_vector<typename map_type::ref_type<cuco::insert_tag>> map_refs;
  map_refs.reserve(num_maps);
  maps.reserve(num_maps);

for (std::size_t i = 0; i < num_maps; ++i) {
    // create a map object and immediately move it into the vector
    maps[i] = std::move(map_type(capacity, cuco::empty_key{empty_key_sentinel}, cuco::empty_value{empty_value_sentinel}));
    // get a non-owning reference to the map which can be consumed in the kernel
    map_refs[i] = maps[i].ref(cuco::insert);
  }

insert_into_maps<<<512, 512>>>(thrust::raw_pointer_cast(map_refs.data()), num_maps);
}