-
Notifications
You must be signed in to change notification settings - Fork 197
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[FEA] Add support for select_k
on CSR matrix
#2140
Conversation
- This PR is one part of the feature of rapidsai#1969 - Add the API of 'select_k' accepting CSR as input - Add the API of 'segmented_copy' Authors: - James Rong (https://github.com/rhdong) Approvers: - Ben Frederickson (https://github.com/benfred) - Micka (https://github.com/lowener) - Corey J. Nolet (https://github.com/cjnolet)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the PR!
I only have a concern regarding the use of the memory resources and device properties. In general, we should use raft "workspace resource" for all temporary allocations in raft algorithms. This gives a greater control, the explicit interface to query the available memory, and slightly better performance in some edge cases.
For your reference, there's another PR that fixes this for all other parts of matrix::select_k
: #1786
Thank you for pointing this out; let me fix it! |
@rhdong can you fix the merge conflicts in this PR please? |
…into rhdong/select_k_csr
Hi @achirkin, @benfred, @cjnolet , when you have a moment, could you please help review this PR? Thank you! Here is the benchmark on A100 with 80GB PCIE.
Original SelectK:
After-optimized SelectK:
SelectK_CSR(Sparisity 0.1/0.2/0.5):
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks, it's nice to see much of the select-k code being reused with minimal modifications! Here's my input; mostly minor things, but please check the alignment of the buffers in the radix select.
Hi @achirkin, Thank you for these valuable review comments, I will fix them as soon as possible! |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for addressing the comments, LGTM
@@ -44,6 +45,16 @@ void select_k(raft::resources const& handle, | |||
bool sorted = false, | |||
SelectAlgo algo = SelectAlgo::kAuto, | |||
const IdxT* len_i = nullptr) RAFT_EXPLICIT; | |||
|
|||
template <typename T, typename IdxT> | |||
void select_k(raft::resources const& handle, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sorry, @rhdong, I do have a nitpick here, now that I finally have time to go through this PR. I'm okay doing this change in an immediate follow-on, but the sparse k-selection should be in raft/sparse/matrix, not in the dense APIs.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hi @cjnolet , thank you for the comment! I just moved the public API to the raft::sparse
scope. Considering many shared source codes, I temporarily didn't move more code to the sparse directory like detail::select_k
; if you prefer more, please let me know.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sorry @rhdong please move the impl details over to sparse also. This is introducing too much confusion having them not be co-located. Also, where possible we should have implementation details calling into public apis for different namespaces.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sorry @rhdong please move the impl details over to sparse also. This is introducing too much confusion having them not be co-located.
Sure, I will.
…into rhdong/select_k_csr
/merge |
Authors:
Approvers: