[ENHANCEMENT]: Add `for_each` device API

[sleeepyjack]

Is your feature request related to a problem? Please describe.

There's currently no way of iterating over all matching values for a given key from the device.
The current retrieve API isn't exposed to the public since it is highly optimized for bulk retrieve and thus not well-suited for general use.

Describe the solution you'd like

Implement a device function that allows for iterating over all values found and then process them according to some user-defined callback function:

template <class CG, class ProbeKey, class Callback>
void for_each(CG tile, ProbeKey key, Callback f) {...}

During probing for key, if a thread in tile hits a matching slot, it will pass a non-incrementable iterator to said slot to the callback which then handles the downstream processing.

This API can also be used to implement the aforementioned optimized retrieve.

Describe alternatives you've considered

for_each handles all of the probing logic internally and just calls the callback with the found values.
The counter design to this would consist of a "heavy" iterator which is returned from a new retrieve device functionThe iterator will handle the probing internally, returningend()once there are no more matching values for the givenkey`.
I have some concerns with this design:

• If we use window loads, i.e., load multiple slots per thread at once, then the iterator has to store this window internally
• The iterator needs to be kept in the same CG context. If we can't ensure that the iterator is incremented in a CG context we might loose some of the values. There's no way we can ensure this never happens with this design. The for_each function, however, ensures that probing is always performed in the same CG context

Additional context

WarpCore had a for_each function although my design back in the day wasn't perfect: https://github.com/sleeepyjack/warpcore/blob/ba5a6a84a8797464e2ce8dd03b269907ee70e7c2/include/warpcore/multi_value_hash_table.cuh#L607

[PointKernel]

The for_each feature makes sense for multiset/multimap and seems invalid for single set/map data structures.

[michaelfliu]

This feature would be very useful for an application I am working on using cuco::static_multimap. My problem requires a mapping from a 4-tuple of positive integers to many positive integers.

Aside: Due to an compilation error in retrieve, I have had to split out the first two/three elements of the key. This was done by creating a vector<vector<vector<cuco::static_multimap<int, int>>>> and using the vector indices to represent the first three elements of the key. The error comes from static_assert(sizeof(Key) == sizeof(cuda::atomic<Key>)); in detail/static_multimap/device_view_impl.inl when Key is thrust::tuple<int, int, int>

They key and associated values are combined to index into a high-dimensional array, producing several vectors of numbers of equal length. I would like to be able to do arbitrary operations on these vectors, such as taking their average, element-wise min, max, etc.

For now, my design idea is:

1. Use the existing cuco::for_each to get the number of values associated with each key
2. Allocate (one large, or many small) device-side arrays to store the values to be retrieved
3. Allocate device memory to hold the result of my operation (e.g. averages, mins, maxes, etc.)
4. Use cuco::for_each to retrieve the values for each key, then perform the operation

With the proposed approach, it seems like I could do the following:

1. Allocate device memory to hold the result of my operation (e.g. averages, mins, maxes, etc.)
2. Use the new cuco::static_multimap::for_each to do the operation

[PointKernel]

@michaelfliu Thanks for the feedback! We will look into it and keep you posted.

In the meantime, we highly recommend to check the mapping table example and there will be two major benefits for your case:

• Get rid of build time error: the key type will be index type thus the static_assert won't fail.
• Improve performance: inserting keys/payloads larger than 8 bytes, e.g. thrust::tuple<int, int, int>, is extremely expensive since lock tables are needed.