Will cuco::static_multimap get a for_each API, or is there a way to get similar functionality?

[michaelfliu]

#506 added a for_each API for static_multiset. Is there any thoughts on adding it for static_multimap?

Alternatively, is there a way to get similar behaviour with the APIs that are currently available for static_multimap?

Specifically, suppose we have a cuco::static_multimap<int, int> mapping integer keys to an arbitrary number of integer values. In my application, the values are indices into another data structure, but to simplify the problem, we can assume they are just positive numbers. For example, we might have the following mappings:

3  -> [1, 2, 4, 5]
10 -> [2, 4]
16 -> [7, 9]

Given a list of keys (e.g. [3, 16]), I would like to compute the average of the values. In this case, I would like a list of two elements

result = {
    3,  // (1 + 2 + 4 + 5) / 4
    8   // (7 + 9) / 2
}

This could be done using a for_each API by first zero-initializing a result array and a count array with the same number of elements as there are keys, then accumulating the sum and count into each of them respectively, and finally dividing element-wise.

With the existing APIs, would the approach be to first count the total number of values associated with the keys (4+2=6 in the above example), then allocating a thrust::device_vector<cuco::pair<int, int>> of size 6 into which we retrieve all of the mapped pairs? The array might look like

retrieved_pairs = {
    {3, 1},
    {3, 4},
    {16, 7},
    {3, 2},
    {16, 9},
    {3, 5}
}

Afterwards, the same sum and count accumulators could be used to compute the required average. Does this method seem reasonable?

[PointKernel]

Does this method seem reasonable?

Yes, absolutely. Here's an example that may help simplify the work: https://godbolt.org/z/Weoj8WKba. The build error will be gone once #529 gets merged.

I was trying to fuse the "retrieve" and "count" operations into one kernel for better performance but then realized some desired APIs like view::initial_slot are not publicly available. Sorry for the inconvenience.

In the meantime, we are deprecating the legacy multimap and the new one will be shipping out in the next month or two. With the new multimap, features like for_each are exposed to the public and writing your custom kernel doing sum reduction will become possible. Stay tuned.

[michaelfliu]

Thank you! I hadn't seen the thrust::reduce_by_key function before, that looks like exactly what I need!

[sleeepyjack]

Note that we also have a new feature called static_map::insert_or_apply which also performs a reduce_by_key operation and shows better performance over thrust::reduce_by_key in most cases. It can be implemented through a custom reduction operator that performs the sum reduction and keeps track of the count as described in your above method: https://godbolt.org/z/4GTesWGjb