[FEA] Add support for `select_k` on CSR matrix #2140

rhdong · 2024-01-31T16:12:50Z

This PR is one part of the feature of [FEA] Pre-filtered brute-force KNN #1969
Add the API of 'select_k' accepting CSR as input
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)

- 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)

achirkin

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

cpp/include/raft/matrix/detail/matrix.cuh

cpp/include/raft/matrix/detail/select_k-inl.cuh

rhdong · 2024-01-31T18:19:45Z

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!

cjnolet · 2024-03-04T17:27:32Z

@rhdong can you fix the merge conflicts in this PR please?

…_k_csr

…into rhdong/select_k_csr

…_k_csr

rhdong · 2024-04-01T21:13:33Z

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.

Try to reuse the CUDA code of select_k as much as possible without hurting its performance.

Original SelectK:

-------------------------------------------------------------------------------------------------------
Benchmark                                                             Time             CPU   Iterations
-------------------------------------------------------------------------------------------------------
SelectK/float/uint32_t/kAuto/8/manual_time                        0.629 ms        0.712 ms         1113 20000#500#256
SelectK/float/uint32_t/kAuto/17/manual_time                       0.296 ms        0.380 ms         2360 1000#10000#256
SelectK/float/uint32_t/kAuto/26/manual_time                       0.400 ms        0.483 ms         1751 100#100000#256
SelectK/float/uint32_t/kAuto/35/manual_time                       0.487 ms        0.570 ms         1439 10#1000000#256

After-optimized SelectK:

-------------------------------------------------------------------------------------------------------
Benchmark                                                             Time             CPU   Iterations
-------------------------------------------------------------------------------------------------------
SelectK/float/uint32_t/kAuto/8/manual_time                        0.623 ms        0.702 ms         1122 20000#500#256
SelectK/float/uint32_t/kAuto/17/manual_time                       0.322 ms        0.403 ms         2177 1000#10000#256
SelectK/float/uint32_t/kAuto/26/manual_time                       0.354 ms        0.436 ms         1978 100#100000#256
SelectK/float/uint32_t/kAuto/35/manual_time                       0.428 ms        0.508 ms         1634 10#1000000#256

SelectK_CSR(Sparisity 0.1/0.2/0.5):

-------------------------------------------------------------------------------------------------------
Benchmark                                                             Time             CPU   Iterations
-------------------------------------------------------------------------------------------------------
SelectKCsrTest<float, uint32_t>/8/manual_time                     0.104 ms        0.186 ms         6692 20000#500#256#0.1
SelectKCsrTest<float, uint32_t>/17/manual_time                    0.115 ms        0.198 ms         6125 1000#10000#256#0.1
SelectKCsrTest<float, uint32_t>/26/manual_time                    0.135 ms        0.219 ms         5178 100#100000#256#0.1
SelectKCsrTest<float, uint32_t>/35/manual_time                    0.214 ms        0.297 ms         3232 10#1000000#256#0.1

SelectKCsrTest<float, uint32_t>/78/manual_time                    0.175 ms        0.257 ms         4006 20000#500#256#0.2
SelectKCsrTest<float, uint32_t>/87/manual_time                    0.161 ms        0.245 ms         4361 1000#10000#256#0.2
SelectKCsrTest<float, uint32_t>/96/manual_time                    0.194 ms        0.277 ms         3616 100#100000#256#0.2
SelectKCsrTest<float, uint32_t>/105/manual_time                   0.268 ms        0.352 ms         2609 10#1000000#256#0.2

SelectKCsrTest<float, uint32_t>/148/manual_time                   0.370 ms        0.452 ms         1896 20000#500#256#0.5
SelectKCsrTest<float, uint32_t>/157/manual_time                   0.240 ms        0.323 ms         2922 1000#10000#256#0.5
SelectKCsrTest<float, uint32_t>/166/manual_time                   0.276 ms        0.359 ms         2541 100#100000#256#0.5
SelectKCsrTest<float, uint32_t>/175/manual_time                   0.353 ms        0.436 ms         2005 10#1000000#256#0.5

achirkin

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.

cpp/include/raft/matrix/detail/select_k-inl.cuh

cpp/test/matrix/select_k_csr.cu

cpp/include/raft/matrix/copy.cuh

cpp/include/raft/matrix/detail/matrix.cuh

cpp/include/raft/matrix/detail/select_radix.cuh

rhdong · 2024-04-02T15:48:00Z

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!

achirkin

Thanks for addressing the comments, LGTM

cjnolet · 2024-04-03T20:08:32Z

cpp/include/raft/matrix/detail/select_k-ext.cuh

@@ -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,


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.

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.

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.

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

cjnolet · 2024-04-08T16:35:36Z

/merge

rhdong requested review from a team as code owners January 31, 2024 16:12

github-actions bot added cpp CMake labels Jan 31, 2024

rhdong requested review from benfred, cjnolet and lowener January 31, 2024 16:13

rhdong added enhancement New feature or request feature request New feature or request non-breaking Non-breaking change 3 - Ready for Review labels Jan 31, 2024

achirkin requested changes Jan 31, 2024

View reviewed changes

Merge branch 'branch-24.04' into rhdong/select_k_csr

7cd874f

cjnolet assigned rhdong Mar 4, 2024

rhdong removed the 3 - Ready for Review label Mar 4, 2024

rhdong added 11 commits March 4, 2024 10:06

Merge remote-tracking branch 'origin/branch-24.04' into rhdong/select…

8550849

…_k_csr

Merge remote-tracking branch 'origin/branch-24.04' into rhdong/select…

3cf7e07

…_k_csr

add more comments on the select_k API

435286a

remove mr argument

977ccee

Merge remote-tracking branch 'origin/branch-24.04' into rhdong/select…

2108c42

…_k_csr

Merge branch 'branch-24.04' into rhdong/select_k_csr

60cfbf3

Merge remote-tracking branch 'origin/branch-24.04' into rhdong/select…

27d1645

…_k_csr

fix format issue

bc544e3

Merge branch 'rhdong/select_k_csr' of https://github.com/rhdong/raft …

a3051d2

…into rhdong/select_k_csr

Merge remote-tracking branch 'origin/branch-24.04' into rhdong/select…

8b9181b

…_k_csr

Optimizing the performance by reusing the dense select_k

558b69e

rhdong changed the base branch from branch-24.04 to branch-24.06 March 30, 2024 08:30

Merge branch 'branch-24.06' into rhdong/select_k_csr

392d050

rhdong added 3 - Ready for Review Public API labels Mar 30, 2024

rhdong requested a review from achirkin March 30, 2024 08:38

rhdong added 2 commits March 30, 2024 10:17

add algo to new select_k

5abbd4e

add more test cases for coverage

ae70762

achirkin requested changes Apr 2, 2024

View reviewed changes

rhdong added 2 commits April 2, 2024 15:37

Response to the review comments.

db24068

fix CI by removing segmented_copy totally.

96ded4d

achirkin approved these changes Apr 3, 2024

View reviewed changes

cjnolet reviewed Apr 3, 2024

View reviewed changes

rhdong added 5 commits April 4, 2024 11:25

move public API to naming scope of sparse

9e8ae31

Merge branch 'branch-24.06' into rhdong/select_k_csr

5924bbb

add back the missed '@defgroup select_k'

b53cee3

Merge branch 'rhdong/select_k_csr' of https://github.com/rhdong/raft …

48cc460

…into rhdong/select_k_csr

Move all of the impl over to sparse name scope

04a8bb4

rhdong added 5 - Ready to Merge and removed 3 - Ready for Review labels Apr 8, 2024

cjnolet approved these changes Apr 8, 2024

View reviewed changes

rapids-bot bot merged commit 4a20d03 into rapidsai:branch-24.06 Apr 8, 2024
73 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[FEA] Add support for `select_k` on CSR matrix #2140

[FEA] Add support for `select_k` on CSR matrix #2140

rhdong commented Jan 31, 2024 •

edited

Loading

achirkin left a comment

rhdong commented Jan 31, 2024

cjnolet commented Mar 4, 2024

rhdong commented Apr 1, 2024 •

edited

Loading

achirkin left a comment

rhdong commented Apr 2, 2024

achirkin left a comment

cjnolet Apr 3, 2024

rhdong Apr 4, 2024

cjnolet Apr 5, 2024 •

edited

Loading

rhdong Apr 5, 2024 •

edited

Loading

cjnolet commented Apr 8, 2024

[FEA] Add support for select_k on CSR matrix #2140

[FEA] Add support for select_k on CSR matrix #2140

Conversation

rhdong commented Jan 31, 2024 • edited Loading

achirkin left a comment

Choose a reason for hiding this comment

rhdong commented Jan 31, 2024

cjnolet commented Mar 4, 2024

rhdong commented Apr 1, 2024 • edited Loading

Original SelectK:

After-optimized SelectK:

SelectK_CSR(Sparisity 0.1/0.2/0.5):

achirkin left a comment

Choose a reason for hiding this comment

rhdong commented Apr 2, 2024

achirkin left a comment

Choose a reason for hiding this comment

cjnolet Apr 3, 2024

Choose a reason for hiding this comment

rhdong Apr 4, 2024

Choose a reason for hiding this comment

cjnolet Apr 5, 2024 • edited Loading

Choose a reason for hiding this comment

rhdong Apr 5, 2024 • edited Loading

Choose a reason for hiding this comment

cjnolet commented Apr 8, 2024

[FEA] Add support for `select_k` on CSR matrix #2140

[FEA] Add support for `select_k` on CSR matrix #2140

rhdong commented Jan 31, 2024 •

edited

Loading

rhdong commented Apr 1, 2024 •

edited

Loading

cjnolet Apr 5, 2024 •

edited

Loading

rhdong Apr 5, 2024 •

edited

Loading