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

Add fused cosine 1-NN cutlass based kernel #2125

Merged
merged 23 commits into from
Mar 19, 2024
Merged

Add fused cosine 1-NN cutlass based kernel #2125

merged 23 commits into from
Mar 19, 2024

Conversation

mdoijade
Copy link
Contributor

@mdoijade mdoijade commented Jan 24, 2024 •
edited
Loading

  • Adds cosine 1-NN cutlass based kernel for SM 8.0 or higher using tensor cores.
  • based on 3x TF32
  • unifies the fusedDistanceNN kernels for L2/cosine.
  • expose this API in pylibraft as fused_distance_nn_arg_min supporting cosine & L2 distance metrics.

@mdoijade mdoijade requested review from a team as code owners January 24, 2024 14:46
@copy-pr-bot copy-pr-bot
Copy link

copy-pr-bot bot commented Jan 24, 2024

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@cjnolet cjnolet added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Jan 24, 2024
@tfeher
Copy link
Contributor

tfeher commented Jan 24, 2024

/ok to test

1 similar comment
@tfeher
Copy link
Contributor

tfeher commented Jan 25, 2024

/ok to test

bdice
bdice reviewed Jan 31, 2024
View reviewed changes
Copy link
Contributor

@bdice bdice left a comment

Choose a reason for hiding this comment

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

Kernels must have internal linkage. See #1898, #1935.

cpp/include/raft/distance/detail/fused_distance_nn/cutlass_base.cuh Outdated Show resolved Hide resolved
cpp/include/raft/distance/detail/fused_distance_nn/simt_kernel.cuh Outdated Show resolved Hide resolved
cpp/include/raft/distance/detail/fused_distance_nn/helper_structs.cuh Outdated Show resolved Hide resolved
@mdoijade
Add fused cosine 1-NN kernel and unify the fused distance 1-NN kernels
1417a2e 
fix doc issue in fused_distance_nn runtime API
tfeher
tfeher requested changes Feb 15, 2024
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 @mdoijade for this PR, it is great to have the fused cosine distance kernels to further accelerate k-means clustering!

It would be nice if we could further improve code reuse between L2 and Cosine variants of fusedNN (decrease duplication of boiler plate code), but I see that we have some discrepancy in the required parameters, and this could negatively affect readibility.

Otherwise the PR looks good in general. Here are my comments

cpp/include/raft/distance/detail/fused_distance_nn/persistent_gemm.h Outdated Show resolved Hide resolved
cpp/include/raft/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h Outdated Show resolved Hide resolved
cpp/include/raft/distance/detail/fused_distance_nn/predicated_tile_iterator_reduced_vec.h Outdated Show resolved Hide resolved
cpp/include/raft/distance/detail/fused_distance_nn/simt_kernel.cuh Show resolved Hide resolved
cpp/include/raft/distance/detail/fused_distance_nn/helper_structs.cuh Show resolved Hide resolved
cpp/src/distance/fused_distance_nn.cu Outdated Show resolved Hide resolved
cpp/include/raft/distance/fused_distance_nn_helpers.cuh Outdated Show resolved Hide resolved
cpp/include/raft/distance/fused_distance_nn.cuh Outdated Show resolved Hide resolved
cpp/include/raft/distance/detail/fused_distance_nn/fused_cosine_nn.cuh Outdated Show resolved Hide resolved
cpp/include/raft/distance/detail/fused_distance_nn.cuh Outdated Show resolved Hide resolved
@mdoijade mdoijade requested review from a team as code owners February 16, 2024 13:19
@github-actions github-actions bot removed the python label Feb 16, 2024
tfeher
tfeher approved these changes Feb 27, 2024
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 @mdoijade for addressing the issues! The code looks good to me.

Please still have a look at the copyright start year in some of the new files.

@mdoijade
Copy link
Contributor Author

mdoijade commented Mar 1, 2024

the pylibraft test passes but an unrelated cpp test is failing LINALG_TEST certainly seems not due to this PR.

@tfeher
Copy link
Contributor

tfeher commented Mar 15, 2024

/rerun tests

@tfeher
Copy link
Contributor

tfeher commented Mar 18, 2024

/merge

benfred
benfred reviewed Mar 18, 2024
View reviewed changes
Copy link
Member

@benfred benfred left a comment

Choose a reason for hiding this comment

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

looks good! one minor comment:

python/pylibraft/pylibraft/distance/fused_distance_nn.pyx Outdated Show resolved Hide resolved
@rapids-bot rapids-bot bot merged commit 413e34e into rapidsai:branch-24.04 Mar 19, 2024
76 of 77 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@bdice bdice bdice left review comments

@tfeher tfeher tfeher approved these changes

@benfred benfred benfred approved these changes

Assignees

@mdoijade mdoijade

Labels
CMake cpp improvement Improvement / enhancement to an existing function non-breaking Non-breaking change python
Projects
VS/ML/DM Primitives Release Board
Status: Done
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
5 participants
@mdoijade @tfeher @benfred @bdice @cjnolet