Shared-memory-cached kernel for reduce_cols_by_key to limit atomic conflicts
#1050
Conversation
Nyrio
added
3 - Ready for Review
improvement
|
Comparison of the performance of the direct vs cached kernels.
|
| */ | ||
| template <typename ElementType, typename KeyType = ElementType, typename IndexType = std::uint32_t> | ||
| void reduce_cols_by_key( | ||
| const raft::handle_t& handle, | ||
| raft::device_matrix_view<const ElementType, IndexType, raft::row_major> data, | ||
| raft::device_vector_view<const KeyType, IndexType> keys, | ||
| raft::device_matrix_view<ElementType, IndexType, raft::row_major> out, | ||
| IndexType nkeys) | ||
| IndexType nkeys, |
There was a problem hiding this comment.
Not part of this PR, but I wonder if there could someday be an overload of this function that would compute this value when not specified.
There was a problem hiding this comment.
Sure, but what kind of use cases do you have in mind? Since this primitive works with keys in the range [0, nkeys[ it's fair to assume nkeys will typically be known. For a more generic primitive working with arbitrary keys, we would first need to map them to such a range, and we can compute nkeys in the process.
There was a problem hiding this comment.
Unless you mean that instead of passing nkeys we should get it from the dimensions of out, which actually makes a lot of sense.
There was a problem hiding this comment.
Yeah, I guess the latter option makes sense. I was considering doing something similar in the knn APIs, since the shape of the output matrices already tell us k.
There was a problem hiding this comment.
I see the following options, let me know if you have a preference:
- Give
nkeysa default value and ignore it, using output dims instead (might break compilation due to unused variable warning?) - Give
nkeysa default value and if not overridden by the user, get it from the output dims. - Remove
nkeysarg (potentially breaking but I doubt anyone is using this prim outside the raft codebase?)
There was a problem hiding this comment.
The second option looks good to me.
There was a problem hiding this comment.
@Nyrio do you want me to go ahead and merge this PR in the meantime?
There was a problem hiding this comment.
Since we aren't in a hurry, I just did the change in this PR. I'd rather put nkeys in last position so we can specify reset_sums and ignore nkeys, but in the interest of not breaking the API I can't do that. If we don't mind breaking the API, I can simply remove the arg and always infer.
| */ | ||
| template <typename ElementType, typename KeyType = ElementType, typename IndexType = std::uint32_t> | ||
| void reduce_cols_by_key( | ||
| const raft::handle_t& handle, | ||
| raft::device_matrix_view<const ElementType, IndexType, raft::row_major> data, | ||
| raft::device_vector_view<const KeyType, IndexType> keys, | ||
| raft::device_matrix_view<ElementType, IndexType, raft::row_major> out, | ||
| IndexType nkeys) | ||
| IndexType nkeys, |
There was a problem hiding this comment.
The second option looks good to me.
|
@gpucibot merge |
reduce_cols_by_keyis used in k-means to compute the total weight in each cluster. When the number of clusters is really small, atomic conflicts make the reduction effectively sequential. This can be avoided by using a shared-memory cache.This PR adds a cached kernel and some heuristics to decide when to use it, resulting in large speedups for
reduce_cols_by_keysand in turn ~15% speedups in some cases for k-means.