Is there any example of static_multimap executing retrieve() in a custom kernel?

[Po1sson]

Is there any example of static_multimap executing retrieve() in a custom kernel? I found it is quite different from insert, involving buffering. I want to know what does "// Final flush of output buffer" mean in kernels.cuh? Where are the results stored before this?

[PointKernel]

Related to #528.

The current device retrieve is designed to write all matches for a given key efficiently. The fact that you need to write a custom kernel implies that for_each is probably a better fit as opposed to retrieve. We will put up the new static_multimap having this feature in the next few weeks. Will keep you posted.

[Po1sson]

Thank you, your answer made me suddenly enlightened. When I saw #528 before, I didn’t realize that this was what I needed.

[avithemad]

@PointKernel : I had a similar question. The use case is that I need to probe multiple static_multimap within a single kernel.
So, the map has a <key, index> pair which indexes into another table. I need to retrieve these indices from the map, and index into another array to get the actual value (so here in the example, the part_map.retrieve(lo_partkey[tid]) should give the indices I need to index into p_brand table and use this to build the result.
Im able to use the contains within a kernel to check its existence as follows:

template <typename Map>
__global__ void probe(Map part_map,
                      Map date_map,
                      Map supp_map,
                      int* lo_partkey,
                      int* lo_orderdate,
                      int* lo_suppkey,
                      int* lo_revenue,
                      int* d_year,
                      int* p_brand,
                      int* res,
                      int lo_size)
{
  int tid = (threadIdx.x + blockIdx.x * blockDim.x) / TILE_SIZE;
  if (tid >= lo_size) return;
  auto this_thread = cg::tiled_partition<TILE_SIZE>(cg::this_thread_block());

  if (!part_map.contains(this_thread, lo_partkey[tid]) ||
      !supp_map.contains(this_thread, lo_suppkey[tid]))
    return;

  res[tid] = 1;
}

But I need to be able to retrieve the keys from the map in order to build the result, therefore needed the retrieve api, is there a better way to do this?

[PointKernel]

@avithemad contains only checks for the presence of a match in the multimap without iterating through all matches. If you're using cuco::experimental::static_multimap, probing_scheme and storage_ref are publicly accessible, allowing you to fully customize operations within the container.

[avithemad]

@PointKernel : Thanks for the reply. Yeah for now I was checking just the count using contains, since I do not know how to probe. I am having trouble on how to use these apis, Is there some example implementation you could share, I would be grateful.

Can I use the initial_slot from the non experimental instead to probe in the device kernel?

[sleeepyjack]

@avithemad would it be possible for you to switch to the new cuco::static_multimap (formerly cuco::experimental::static_multimap edit: the new multimap is still named cuco::experimental::static_multimap? I will put up a PR that introduces the for_each APIs ASAP. In the meantime, you could start switching to the new multimap implementation in your codebase. I will ping you once the PR is up so you can give it a try.

[sleeepyjack]

@avithemad #599 adds the static_multimap_ref::for_each device API. Feel free to checkout the PR branch and see if it works for your problem.

[avithemad]

Thank you, I will try it out.