Skip to content
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.

Sign up for GitHub

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

Jump to bottom

Integrate KNN implementation: ivf-pq #789

Merged
merged 148 commits into from
Sep 30, 2022
Merged

Integrate KNN implementation: ivf-pq #789

merged 148 commits into from
Sep 30, 2022

Conversation

achirkin
Copy link
Contributor

Integrate a new KNN implementation.

achirkin added 20 commits June 24, 2022 14:37
@achirkin
Initial commit: copied everything into a single file
4d7b699
@achirkin
Merge branch 'branch-22.08' into fea-knn-ivf-pq
306d4bb
@achirkin
Add missing functions and fix small compile errors
ebc20be
@achirkin
Added tests (still failing)
4f71b09
@achirkin
WIP replacing chunks of code with raft's helpers
1d6402f
@achirkin
Reset the expected cuda error to cudaSuccess to not crumble following…
f12e5da 
… cuda calls
@achirkin
Replace cuann return codes with raft exceptions
12fad92
@achirkin
replace cublas calls with raft wrappers
0bb702f
@achirkin
Replace _cuann_memset with utils::memzero
e1cae88
@achirkin
Use raft logging in most of the places and wrap all cuda calls into R…
5b2c799 
…AFT_CUDA_TRY
@achirkin
Use helper accumulate_into_selected
4f64cd3
@achirkin
Use helper copy_selected
0a7324b
@achirkin
Wrap remaining cuda calls into RAFT_CUDA_TRY and replace asserts with…
4d11178 
… RAFT_EXPECTS
@achirkin
WIP: templatizing the api
267b35e
@achirkin
Add more templates and use adjust_centers
2541250
@achirkin
Put index as a member of descriptor
67bffa4
@achirkin
Add the index extending operation
f951c52
@achirkin
Replace cuann similarity with raft distance metric
5822844
@achirkin
Use raft logging levels instead of CUANN_DEBUG macro
62d32a7
@achirkin
Made the API scaffold.
48c8702
@achirkin achirkin added feature request New feature or request non-breaking Non-breaking change 2 - In Progress Currenty a work in progress labels Aug 15, 2022
@achirkin achirkin self-assigned this Aug 15, 2022
@achirkin
Copy link
Contributor Author

achirkin commented Aug 15, 2022 •
edited
Loading

Current state

I've made the public API that mirrors ivf_flat. The public functions in raft/spatial/knn/ivf_pq.cuh and structures (raft/spatial/knn/ivf_pq_types.hpp) shouldn't change much anymore. Functionality-wise: it works; performance (especially that of training) is to be investigated and improved in a follow-up PR.

Major TODOs

  • ivf_pq::extend function ignores indices at the moment - fix that
  • Fix training/search for the InnerProduct distance metric (cuann version works only under assumption of normalized data and queries)
  • Propagate template IdxT for the index type (original implementation keeps indices as uint32_t, but outputs them as uint64_t)
  • Expand tests and check correctness for the implementation-specific cases (float16 intermediate type, combinations of pq_bits and pq_dims, etc.)
  • Remove cuannIvfPqDescriptor and its void* index member and merge all their content into ivf_pq::index; switch to mdarrays for storing device data
  • Replace all allocations with rmm-managed buffers/vectors
  • Reuse/move missing helpers into detail/ann_utils.cuh
  • Replace BlockTopk with raft's recently integrated topk structures.
  • Consolidate kmeans-related code and try to reuse detail/ann_kmeans_balanced.cuh as much as possible
  • Eventually remove detail/ivf_pq_legacy.cuh
  • Include the new implementation into the synthetic benchmarks at bench/spatial/knn.cu
  • The training takes way too long. Visual inspection of nsys timeline suggests, utils::argmin_along_rows_kernels is a major bottleneck and should be optimized; plus, modification of the trainset on CPU in the PER_SUBSPACE case takes one third of all training time and should be moved to GPU (mod_trainset[] = transpose( rotate(trainset[]) - clusterRotCenters[] )).

@achirkin
Copy link
Contributor Author

achirkin commented Aug 15, 2022 •
edited
Loading

NB: this PR depends on #788 (changes in kmeans api).

@achirkin achirkin requested a review from tfeher September 30, 2022 14:35
@achirkin achirkin requested a review from cjnolet September 30, 2022 14:46
tfeher
tfeher approved these changes Sep 30, 2022
View reviewed changes
Copy link
Contributor

@tfeher tfeher left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks Artem for addressing the issues! The PR looks good to me!

@cjnolet
Copy link
Member

cjnolet commented Sep 30, 2022

@gpucibot merge

@rapids-bot rapids-bot bot merged commit 8a31ae6 into rapidsai:branch-22.10 Sep 30, 2022
@achirkin achirkin deleted the fea-knn-ivf-pq branch October 1, 2022 05:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@Nyrio Nyrio Nyrio left review comments

@tfeher tfeher tfeher approved these changes

@cjnolet cjnolet cjnolet approved these changes

Assignees

@achirkin achirkin

Labels
2 - In Progress Currenty a work in progress CMake cpp feature request New feature or request non-breaking Non-breaking change
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@achirkin @cjnolet @tfeher @Nyrio